mz_transform/
predicate_pushdown.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//! Pushes predicates down through other operators.
11//!
12//! This action generally improves the quality of the query, in that selective per-record
13//! filters reduce the volume of data before they arrive at more expensive operators.
14//!
15//!
16//! The one time when this action might not improve the quality of a query is
17//! if a filter gets pushed down on an arrangement because that blocks arrangement
18//! reuse. It assumed that actions that need an arrangement are responsible for
19//! lifting filters out of the way.
20//!
21//! Predicate pushdown will not push down literal errors, unless it is certain that
22//! the literal errors will be unconditionally evaluated. For example, the pushdown
23//! will not happen if not all predicates can be pushed down (e.g. reduce and map),
24//! or if we are not certain that the input is non-empty (e.g. join).
25//! Note that this is not addressing the problem in its full generality, because this problem can
26//! occur with any function call that might error (although much more rarely than with literal
27//! errors). See <https://github.com/MaterializeInc/database-issues/issues/4972#issuecomment-1547391011>
28//!
29//! ```rust
30//! use mz_expr::{BinaryFunc, MirRelationExpr, MirScalarExpr};
31//! use mz_ore::id_gen::IdGen;
32//! use mz_repr::{SqlColumnType, Datum, SqlRelationType, SqlScalarType};
33//! use mz_repr::optimize::OptimizerFeatures;
34//! use mz_transform::{typecheck, Transform, TransformCtx};
35//! use mz_transform::dataflow::DataflowMetainfo;
36//!
37//! use mz_transform::predicate_pushdown::PredicatePushdown;
38//!
39//! let input1 = MirRelationExpr::constant(vec![], SqlRelationType::new(vec![
40//!     SqlScalarType::Bool.nullable(false),
41//! ]));
42//! let input2 = MirRelationExpr::constant(vec![], SqlRelationType::new(vec![
43//!     SqlScalarType::Bool.nullable(false),
44//! ]));
45//! let input3 = MirRelationExpr::constant(vec![], SqlRelationType::new(vec![
46//!     SqlScalarType::Bool.nullable(false),
47//! ]));
48//! let join = MirRelationExpr::join(
49//!     vec![input1.clone(), input2.clone(), input3.clone()],
50//!     vec![vec![(0, 0), (2, 0)].into_iter().collect()],
51//! );
52//!
53//! let predicate0 = MirScalarExpr::column(0);
54//! let predicate1 = MirScalarExpr::column(1);
55//! let predicate01 = MirScalarExpr::column(0).call_binary(MirScalarExpr::column(2), BinaryFunc::AddInt64);
56//! let predicate012 = MirScalarExpr::literal_false();
57//!
58//! let mut expr = join.filter(
59//!    vec![
60//!        predicate0.clone(),
61//!        predicate1.clone(),
62//!        predicate01.clone(),
63//!        predicate012.clone(),
64//!    ]);
65//!
66//! let features = OptimizerFeatures::default();
67//! let typecheck_ctx = typecheck::empty_context();
68//! let mut df_meta = DataflowMetainfo::default();
69//! let mut transform_ctx = TransformCtx::local(&features, &typecheck_ctx, &mut df_meta, None, None);
70//!
71//! PredicatePushdown::default().transform(&mut expr, &mut transform_ctx);
72//!
73//! let predicate00 = MirScalarExpr::column(0).call_binary(MirScalarExpr::column(0), BinaryFunc::AddInt64);
74//! let expected_expr = MirRelationExpr::join(
75//!     vec![
76//!         input1.clone().filter(vec![predicate0.clone(), predicate00.clone()]),
77//!         input2.clone().filter(vec![predicate0.clone()]),
78//!         input3.clone().filter(vec![predicate0, predicate00])
79//!     ],
80//!     vec![vec![(0, 0), (2, 0)].into_iter().collect()],
81//! ).filter(vec![predicate012]);
82//! assert_eq!(expected_expr, expr)
83//! ```
84
85use std::collections::{BTreeMap, BTreeSet};
86
87use itertools::Itertools;
88use mz_expr::visit::{Visit, VisitChildren};
89use mz_expr::{
90    AggregateFunc, Id, JoinInputMapper, LocalId, MirRelationExpr, MirScalarExpr, RECURSION_LIMIT,
91    VariadicFunc, func,
92};
93use mz_ore::soft_assert_eq_no_log;
94use mz_ore::stack::{CheckedRecursion, RecursionGuard, RecursionLimitError};
95use mz_repr::{Datum, SqlColumnType, SqlScalarType};
96
97use crate::{TransformCtx, TransformError};
98
99/// Pushes predicates down through other operators.
100#[derive(Debug)]
101pub struct PredicatePushdown {
102    recursion_guard: RecursionGuard,
103}
104
105impl Default for PredicatePushdown {
106    fn default() -> PredicatePushdown {
107        PredicatePushdown {
108            recursion_guard: RecursionGuard::with_limit(RECURSION_LIMIT),
109        }
110    }
111}
112
113impl CheckedRecursion for PredicatePushdown {
114    fn recursion_guard(&self) -> &RecursionGuard {
115        &self.recursion_guard
116    }
117}
118
119impl crate::Transform for PredicatePushdown {
120    fn name(&self) -> &'static str {
121        "PredicatePushdown"
122    }
123
124    #[mz_ore::instrument(
125        target = "optimizer",
126        level = "debug",
127        fields(path.segment = "predicate_pushdown")
128    )]
129    fn actually_perform_transform(
130        &self,
131        relation: &mut MirRelationExpr,
132        _: &mut TransformCtx,
133    ) -> Result<(), TransformError> {
134        let mut empty = BTreeMap::new();
135        let result = self.action(relation, &mut empty);
136        mz_repr::explain::trace_plan(&*relation);
137        result
138    }
139}
140
141impl PredicatePushdown {
142    /// Predicate pushdown
143    ///
144    /// This method looks for opportunities to push predicates toward
145    /// sources of data. Primarily, this is the `Filter` expression,
146    /// and moving its predicates through the operators it contains.
147    ///
148    /// In addition, the method accumulates the intersection of predicates
149    /// applied to each `Get` expression, so that the predicate can
150    /// then be pushed through to a `Let` binding, or to the external
151    /// source of the data if the `Get` binds to another view.
152    pub fn action(
153        &self,
154        relation: &mut MirRelationExpr,
155        get_predicates: &mut BTreeMap<Id, BTreeSet<MirScalarExpr>>,
156    ) -> Result<(), TransformError> {
157        self.checked_recur(|_| {
158            // In the case of Filter or Get we have specific work to do;
159            // otherwise we should recursively descend.
160            match relation {
161                MirRelationExpr::Filter { input, predicates } => {
162                    // Reduce the predicates to determine as best as possible
163                    // whether they are literal errors before working with them.
164                    let input_type = input.typ();
165                    for predicate in predicates.iter_mut() {
166                        predicate.reduce(&input_type.column_types);
167                    }
168
169                    // It can be helpful to know if there are any non-literal errors,
170                    // as this is justification for not pushing down literal errors.
171                    let all_errors = predicates.iter().all(|p| p.is_literal_err());
172                    // Depending on the type of `input` we have different
173                    // logic to apply to consider pushing `predicates` down.
174                    match &mut **input {
175                        MirRelationExpr::Let { body, .. }
176                        | MirRelationExpr::LetRec { body, .. } => {
177                            // Push all predicates to the body.
178                            **body = body
179                                .take_dangerous()
180                                .filter(std::mem::replace(predicates, Vec::new()));
181
182                            self.action(input, get_predicates)?;
183                        }
184                        MirRelationExpr::Get { id, .. } => {
185                            // We can report the predicates upward in `get_predicates`,
186                            // but we are not yet able to delete them from the
187                            // `Filter`.
188                            get_predicates
189                                .entry(*id)
190                                .or_insert_with(|| predicates.iter().cloned().collect())
191                                .retain(|p| predicates.contains(p));
192                        }
193                        MirRelationExpr::Join {
194                            inputs,
195                            equivalences,
196                            ..
197                        } => {
198                            // We want to scan `predicates` for any that can
199                            // 1) become join variable constraints
200                            // 2) apply to individual elements of `inputs`.
201                            // Figuring out the set of predicates that belong to
202                            //    the latter group requires 1) knowing which predicates
203                            //    are in the former group and 2) that the variable
204                            //    constraints be in canonical form.
205                            // Thus, there is a first scan across `predicates` to
206                            //    populate the join variable constraints
207                            //    and a second scan across the remaining predicates
208                            //    to see which ones can become individual elements of
209                            //    `inputs`.
210
211                            let input_mapper = mz_expr::JoinInputMapper::new(inputs);
212
213                            // Predicates not translated into join variable
214                            // constraints. We will attempt to push them at all
215                            // inputs, and failing to
216                            let mut pred_not_translated = Vec::new();
217
218                            for mut predicate in predicates.drain(..) {
219                                use mz_expr::{BinaryFunc, UnaryFunc};
220                                if let MirScalarExpr::CallBinary {
221                                    func: BinaryFunc::Eq,
222                                    expr1,
223                                    expr2,
224                                } = &predicate
225                                {
226                                    // Translate into join variable constraints:
227                                    // 1) `nonliteral1 == nonliteral2` constraints
228                                    // 2) `expr == literal` where `expr` refers to more
229                                    //    than one input.
230                                    let input_count =
231                                        input_mapper.lookup_inputs(&predicate).count();
232                                    if (!expr1.is_literal() && !expr2.is_literal())
233                                        || input_count >= 2
234                                    {
235                                        // `col1 == col2` as a `MirScalarExpr`
236                                        // implies `!isnull(col1)` as well.
237                                        // `col1 == col2` as a join constraint does
238                                        // not have this extra implication.
239                                        // Thus, when translating the
240                                        // `MirScalarExpr` to a join constraint, we
241                                        // need to retain the `!isnull(col1)`
242                                        // information.
243                                        if expr1.typ(&input_type.column_types).nullable {
244                                            pred_not_translated.push(
245                                                expr1
246                                                    .clone()
247                                                    .call_unary(UnaryFunc::IsNull(func::IsNull))
248                                                    .call_unary(UnaryFunc::Not(func::Not)),
249                                            );
250                                        } else if expr2.typ(&input_type.column_types).nullable {
251                                            pred_not_translated.push(
252                                                expr2
253                                                    .clone()
254                                                    .call_unary(UnaryFunc::IsNull(func::IsNull))
255                                                    .call_unary(UnaryFunc::Not(func::Not)),
256                                            );
257                                        }
258                                        equivalences
259                                            .push(vec![(**expr1).clone(), (**expr2).clone()]);
260                                        continue;
261                                    }
262                                } else if let Some((expr1, expr2)) =
263                                    Self::extract_equal_or_both_null(
264                                        &mut predicate,
265                                        &input_type.column_types,
266                                    )
267                                {
268                                    // Also translate into join variable constraints:
269                                    // 3) `((nonliteral1 = nonliteral2) || (nonliteral
270                                    //    is null && nonliteral2 is null))`
271                                    equivalences.push(vec![expr1, expr2]);
272                                    continue;
273                                }
274                                pred_not_translated.push(predicate)
275                            }
276
277                            mz_expr::canonicalize::canonicalize_equivalences(
278                                equivalences,
279                                std::iter::once(&input_type.column_types),
280                            );
281
282                            let (retain, push_downs) = Self::push_filters_through_join(
283                                &input_mapper,
284                                equivalences,
285                                pred_not_translated,
286                            );
287
288                            Self::update_join_inputs_with_push_downs(inputs, push_downs);
289
290                            // Recursively descend on the join
291                            self.action(input, get_predicates)?;
292
293                            // remove all predicates that were pushed down from the current Filter node
294                            *predicates = retain;
295                        }
296                        MirRelationExpr::Reduce {
297                            input: inner,
298                            group_key,
299                            aggregates,
300                            monotonic: _,
301                            expected_group_size: _,
302                        } => {
303                            let mut retain = Vec::new();
304                            let mut push_down = Vec::new();
305                            for predicate in predicates.drain(..) {
306                                // Do not push down literal errors unless it is only errors.
307                                if !predicate.is_literal_err() || all_errors {
308                                    let mut supported = true;
309                                    let mut new_predicate = predicate.clone();
310                                    new_predicate.visit_pre(|e| {
311                                        if let MirScalarExpr::Column(c, _) = e {
312                                            if *c >= group_key.len() {
313                                                supported = false;
314                                            }
315                                        }
316                                    });
317                                    if supported {
318                                        new_predicate.visit_mut_post(&mut |e| {
319                                            if let MirScalarExpr::Column(i, _) = e {
320                                                *e = group_key[*i].clone();
321                                            }
322                                        })?;
323                                        push_down.push(new_predicate);
324                                    } else if let MirScalarExpr::Column(col, _) = &predicate {
325                                        if *col == group_key.len()
326                                            && aggregates.len() == 1
327                                            && aggregates[0].func == AggregateFunc::Any
328                                        {
329                                            push_down.push(aggregates[0].expr.clone());
330                                            aggregates[0].expr = MirScalarExpr::literal_ok(
331                                                Datum::True,
332                                                SqlScalarType::Bool,
333                                            );
334                                        } else {
335                                            retain.push(predicate);
336                                        }
337                                    } else {
338                                        retain.push(predicate);
339                                    }
340                                } else {
341                                    retain.push(predicate);
342                                }
343                            }
344
345                            if !push_down.is_empty() {
346                                *inner = Box::new(inner.take_dangerous().filter(push_down));
347                            }
348                            self.action(inner, get_predicates)?;
349
350                            // remove all predicates that were pushed down from the current Filter node
351                            std::mem::swap(&mut retain, predicates);
352                        }
353                        MirRelationExpr::TopK {
354                            input,
355                            group_key,
356                            order_key: _,
357                            limit,
358                            offset: _,
359                            monotonic: _,
360                            expected_group_size: _,
361                        } => {
362                            let mut retain = Vec::new();
363                            let mut push_down = Vec::new();
364
365                            let mut support = BTreeSet::new();
366                            support.extend(group_key.iter().cloned());
367                            if let Some(limit) = limit {
368                                // Strictly speaking not needed because the
369                                // `limit` support should be a subset of the
370                                // `group_key` support, but we don't want to
371                                // take this for granted here.
372                                limit.support_into(&mut support);
373                            }
374
375                            for predicate in predicates.drain(..) {
376                                // Do not push down literal errors unless it is only errors.
377                                if (!predicate.is_literal_err() || all_errors)
378                                    && predicate.support().is_subset(&support)
379                                {
380                                    push_down.push(predicate);
381                                } else {
382                                    retain.push(predicate);
383                                }
384                            }
385
386                            // remove all predicates that were pushed down from the current Filter node
387                            std::mem::swap(&mut retain, predicates);
388
389                            if !push_down.is_empty() {
390                                *input = Box::new(input.take_dangerous().filter(push_down));
391                            }
392
393                            self.action(input, get_predicates)?;
394                        }
395                        MirRelationExpr::Threshold { input } => {
396                            let predicates = std::mem::take(predicates);
397                            *relation = input.take_dangerous().filter(predicates).threshold();
398                            self.action(relation, get_predicates)?;
399                        }
400                        MirRelationExpr::Project { input, outputs } => {
401                            let predicates = predicates.drain(..).map(|mut predicate| {
402                                predicate.permute(outputs);
403                                predicate
404                            });
405                            *relation = input
406                                .take_dangerous()
407                                .filter(predicates)
408                                .project(outputs.clone());
409
410                            self.action(relation, get_predicates)?;
411                        }
412                        MirRelationExpr::Filter {
413                            input,
414                            predicates: predicates2,
415                        } => {
416                            *relation = input
417                                .take_dangerous()
418                                .filter(predicates.clone().into_iter().chain(predicates2.clone()));
419                            self.action(relation, get_predicates)?;
420                        }
421                        MirRelationExpr::Map { input, scalars } => {
422                            let (retained, pushdown) = Self::push_filters_through_map(
423                                scalars,
424                                predicates,
425                                input.arity(),
426                                all_errors,
427                            )?;
428                            let scalars = std::mem::take(scalars);
429                            let mut result = input.take_dangerous();
430                            if !pushdown.is_empty() {
431                                result = result.filter(pushdown);
432                            }
433                            self.action(&mut result, get_predicates)?;
434                            result = result.map(scalars);
435                            if !retained.is_empty() {
436                                result = result.filter(retained);
437                            }
438                            *relation = result;
439                        }
440                        MirRelationExpr::FlatMap { input, .. } => {
441                            let (mut retained, pushdown) =
442                                Self::push_filters_through_flat_map(predicates, input.arity());
443
444                            // remove all predicates that were pushed down from the current Filter node
445                            std::mem::swap(&mut retained, predicates);
446
447                            if !pushdown.is_empty() {
448                                // put the filter on top of the input
449                                **input = input.take_dangerous().filter(pushdown);
450                            }
451
452                            // ... and keep pushing predicates down
453                            self.action(input, get_predicates)?;
454                        }
455                        MirRelationExpr::Union { base, inputs } => {
456                            let predicates = std::mem::take(predicates);
457                            *base = Box::new(base.take_dangerous().filter(predicates.clone()));
458                            self.action(base, get_predicates)?;
459                            for input in inputs {
460                                *input = input.take_dangerous().filter(predicates.clone());
461                                self.action(input, get_predicates)?;
462                            }
463                        }
464                        MirRelationExpr::Negate { input } => {
465                            // Don't push literal errors past a Negate. The problem is that it's
466                            // hard to appropriately reflect the negation in the error stream:
467                            // - If we don't negate, then errors that should cancel out will not
468                            //   cancel out. For example, see
469                            //   https://github.com/MaterializeInc/database-issues/issues/5691
470                            // - If we negate, then unrelated errors might cancel out. E.g., there
471                            //   might be a division-by-0 in both inputs to an EXCEPT ALL, but
472                            //   on different input data. These shouldn't cancel out.
473                            let (retained, pushdown): (Vec<_>, Vec<_>) = std::mem::take(predicates)
474                                .into_iter()
475                                .partition(|p| p.is_literal_err());
476                            let mut result = input.take_dangerous();
477                            if !pushdown.is_empty() {
478                                result = result.filter(pushdown);
479                            }
480                            self.action(&mut result, get_predicates)?;
481                            result = result.negate();
482                            if !retained.is_empty() {
483                                result = result.filter(retained);
484                            }
485                            *relation = result;
486                        }
487                        x => {
488                            x.try_visit_mut_children(|e| self.action(e, get_predicates))?;
489                        }
490                    }
491
492                    // remove empty filters (junk by-product of the actual transform)
493                    match relation {
494                        MirRelationExpr::Filter { predicates, input } if predicates.is_empty() => {
495                            *relation = input.take_dangerous();
496                        }
497                        _ => {}
498                    }
499
500                    Ok(())
501                }
502                MirRelationExpr::Get { id, .. } => {
503                    // Purge all predicates associated with the id.
504                    get_predicates
505                        .entry(*id)
506                        .or_insert_with(BTreeSet::new)
507                        .clear();
508
509                    Ok(())
510                }
511                MirRelationExpr::Let { id, body, value } => {
512                    // Push predicates and collect intersection at `Get`s.
513                    self.action(body, get_predicates)?;
514
515                    // `get_predicates` should now contain the intersection
516                    // of predicates at each *use* of the binding. If it is
517                    // non-empty, we can move those predicates to the value.
518                    Self::push_into_let_binding(get_predicates, id, value, &mut [body]);
519
520                    // Continue recursively on the value.
521                    self.action(value, get_predicates)
522                }
523                MirRelationExpr::LetRec {
524                    ids,
525                    values,
526                    limits: _,
527                    body,
528                } => {
529                    // Note: This could be extended to be able to do a little more pushdowns, see
530                    // https://github.com/MaterializeInc/database-issues/issues/5336#issuecomment-1477588262
531
532                    // Pre-compute which Ids are used across iterations
533                    let ids_used_across_iterations = MirRelationExpr::recursive_ids(ids, values);
534
535                    // Predicate pushdown within the body
536                    self.action(body, get_predicates)?;
537
538                    // `users` will be the body plus the values of those bindings that we have seen
539                    // so far, while going one-by-one through the list of bindings backwards.
540                    // `users` contains those expressions from which we harvested `get_predicates`,
541                    // and therefore we should attend to all of these expressions when pushing down
542                    // a predicate into a Let binding.
543                    let mut users = vec![&mut **body];
544                    for (id, value) in ids.iter_mut().zip(values).rev() {
545                        // Predicate pushdown from Gets in `users` into the value of a Let binding
546                        //
547                        // For now, we simply always avoid pushing into a Let binding that is
548                        // referenced across iterations to avoid soundness problems and infinite
549                        // pushdowns.
550                        //
551                        // Note that `push_into_let_binding` makes a further check based on
552                        // `get_predicates`: We push a predicate into the value of a binding, only
553                        // if all Gets of this Id have this same predicate on top of them.
554                        if !ids_used_across_iterations.contains(id) {
555                            Self::push_into_let_binding(get_predicates, id, value, &mut users);
556                        }
557
558                        // Predicate pushdown within a binding
559                        self.action(value, get_predicates)?;
560
561                        users.push(value);
562                    }
563
564                    Ok(())
565                }
566                MirRelationExpr::Join {
567                    inputs,
568                    equivalences,
569                    ..
570                } => {
571                    // The goal is to push
572                    //   1) equivalences of the form `expr = <runtime constant>`, where `expr`
573                    //      comes from a single input.
574                    //   2) equivalences of the form `expr1 = expr2`, where both
575                    //      expressions come from the same single input.
576                    let input_types = inputs.iter().map(|i| i.typ()).collect::<Vec<_>>();
577                    mz_expr::canonicalize::canonicalize_equivalences(
578                        equivalences,
579                        input_types.iter().map(|t| &t.column_types),
580                    );
581
582                    let input_mapper = mz_expr::JoinInputMapper::new_from_input_types(&input_types);
583                    // Predicates to push at each input, and to lift out the join.
584                    let mut push_downs = vec![Vec::new(); inputs.len()];
585
586                    for equivalence_pos in 0..equivalences.len() {
587                        // Case 1: there are more than one literal in the
588                        // equivalence class. Because of equivalences have been
589                        // dedupped, this means that everything in the equivalence
590                        // class must be equal to two different literals, so the
591                        // entire relation zeroes out
592                        if equivalences[equivalence_pos]
593                            .iter()
594                            .filter(|expr| expr.is_literal())
595                            .count()
596                            > 1
597                        {
598                            relation.take_safely(Some(relation.typ_with_input_types(&input_types)));
599                            return Ok(());
600                        }
601
602                        let runtime_constants = equivalences[equivalence_pos]
603                            .iter()
604                            .filter(|expr| expr.support().is_empty())
605                            .cloned()
606                            .collect::<Vec<_>>();
607                        if !runtime_constants.is_empty() {
608                            // Case 2: There is at least one runtime constant the equivalence class
609                            let gen_literal_equality_preds = |expr: MirScalarExpr| {
610                                let mut equality_preds = Vec::new();
611                                for constant in runtime_constants.iter() {
612                                    let pred = if constant.is_literal_null() {
613                                        MirScalarExpr::CallUnary {
614                                            func: mz_expr::UnaryFunc::IsNull(func::IsNull),
615                                            expr: Box::new(expr.clone()),
616                                        }
617                                    } else {
618                                        MirScalarExpr::CallBinary {
619                                            func: mz_expr::BinaryFunc::Eq,
620                                            expr1: Box::new(expr.clone()),
621                                            expr2: Box::new(constant.clone()),
622                                        }
623                                    };
624                                    equality_preds.push(pred);
625                                }
626                                equality_preds
627                            };
628
629                            // Find all single input expressions in the equivalence
630                            // class and collect (position within the equivalence class,
631                            // input the expression belongs to, localized version of the
632                            // expression).
633                            let mut single_input_exprs = equivalences[equivalence_pos]
634                                .iter()
635                                .enumerate()
636                                .filter_map(|(pos, e)| {
637                                    let mut inputs = input_mapper.lookup_inputs(e);
638                                    if let Some(input) = inputs.next() {
639                                        if inputs.next().is_none() {
640                                            return Some((
641                                                pos,
642                                                input,
643                                                input_mapper.map_expr_to_local(e.clone()),
644                                            ));
645                                        }
646                                    }
647                                    None
648                                })
649                                .collect::<Vec<_>>();
650
651                            // For every single-input expression `expr`, we can push
652                            // down `expr = <runtime constant>` and remove `expr` from the
653                            // equivalence class.
654                            for (expr_pos, input, expr) in single_input_exprs.drain(..).rev() {
655                                push_downs[input].extend(gen_literal_equality_preds(expr));
656                                equivalences[equivalence_pos].remove(expr_pos);
657                            }
658
659                            // If none of the expressions in the equivalence depend on input
660                            // columns and equality predicates with them are pushed down,
661                            // we can safely remove them from the equivalence.
662                            // TODO: we could probably push equality predicates among the
663                            // remaining constants to all join inputs to prevent any computation
664                            // from happening until the condition is satisfied.
665                            if equivalences[equivalence_pos]
666                                .iter()
667                                .all(|e| e.support().is_empty())
668                                && push_downs.iter().any(|p| !p.is_empty())
669                            {
670                                equivalences[equivalence_pos].clear();
671                            }
672                        } else {
673                            // Case 3: There are no constants in the equivalence
674                            // class. Push a predicate for every pair of expressions
675                            // in the equivalence that either belong to a single
676                            // input or can be localized to a given input through
677                            // the rest of equivalences.
678                            let mut to_remove = Vec::new();
679                            for input in 0..inputs.len() {
680                                // Vector of pairs (position within the equivalence, localized
681                                // expression). The position is None for expressions derived through
682                                // other equivalences.
683                                let localized = equivalences[equivalence_pos]
684                                    .iter()
685                                    .enumerate()
686                                    .filter_map(|(pos, expr)| {
687                                        if let MirScalarExpr::Column(col_pos, _) = &expr {
688                                            let local_col =
689                                                input_mapper.map_column_to_local(*col_pos);
690                                            if input == local_col.1 {
691                                                // TODO(mgree) !!! is it safe to propagate the name here?
692                                                return Some((
693                                                    Some(pos),
694                                                    MirScalarExpr::column(local_col.0),
695                                                ));
696                                            } else {
697                                                return None;
698                                            }
699                                        }
700                                        let mut inputs = input_mapper.lookup_inputs(expr);
701                                        if let Some(single_input) = inputs.next() {
702                                            if input == single_input && inputs.next().is_none() {
703                                                return Some((
704                                                    Some(pos),
705                                                    input_mapper.map_expr_to_local(expr.clone()),
706                                                ));
707                                            }
708                                        }
709                                        // Equivalences not including the current expression
710                                        let mut other_equivalences = equivalences.clone();
711                                        other_equivalences[equivalence_pos].remove(pos);
712                                        let mut localized = expr.clone();
713                                        if input_mapper.try_localize_to_input_with_bound_expr(
714                                            &mut localized,
715                                            input,
716                                            &other_equivalences[..],
717                                        ) {
718                                            Some((None, localized))
719                                        } else {
720                                            None
721                                        }
722                                    })
723                                    .collect::<Vec<_>>();
724
725                                // If there are at least 2 expression in the equivalence that
726                                // can be localized to the same input, push all combinations
727                                // of them to the input.
728                                if localized.len() > 1 {
729                                    for mut pair in
730                                        localized.iter().map(|(_, expr)| expr).combinations(2)
731                                    {
732                                        let expr1 = pair.pop().unwrap();
733                                        let expr2 = pair.pop().unwrap();
734
735                                        push_downs[input].push(
736                                            MirScalarExpr::CallBinary {
737                                                func: mz_expr::BinaryFunc::Eq,
738                                                expr1: Box::new(expr2.clone()),
739                                                expr2: Box::new(expr1.clone()),
740                                            }
741                                            .or(expr2
742                                                .clone()
743                                                .call_is_null()
744                                                .and(expr1.clone().call_is_null())),
745                                        );
746                                    }
747
748                                    if localized.len() == equivalences[equivalence_pos].len() {
749                                        // The equivalence is either a single input one or fully localizable
750                                        // to a single input through other equivalences, so it can be removed
751                                        // completely without introducing any new cross join.
752                                        to_remove.extend(0..equivalences[equivalence_pos].len());
753                                    } else {
754                                        // Leave an expression from this input in the equivalence to avoid
755                                        // cross joins
756                                        to_remove.extend(
757                                            localized.iter().filter_map(|(pos, _)| *pos).skip(1),
758                                        );
759                                    }
760                                }
761                            }
762
763                            // Remove expressions that were pushed down to at least one input
764                            to_remove.sort();
765                            to_remove.dedup();
766                            for pos in to_remove.iter().rev() {
767                                equivalences[equivalence_pos].remove(*pos);
768                            }
769                        };
770                    }
771
772                    mz_expr::canonicalize::canonicalize_equivalences(
773                        equivalences,
774                        input_types.iter().map(|t| &t.column_types),
775                    );
776
777                    Self::update_join_inputs_with_push_downs(inputs, push_downs);
778
779                    // Recursively descend on each of the inputs.
780                    for input in inputs.iter_mut() {
781                        self.action(input, get_predicates)?;
782                    }
783
784                    Ok(())
785                }
786                x => {
787                    // Recursively descend.
788                    x.try_visit_mut_children(|e| self.action(e, get_predicates))
789                }
790            }
791        })
792    }
793
794    fn update_join_inputs_with_push_downs(
795        inputs: &mut Vec<MirRelationExpr>,
796        push_downs: Vec<Vec<MirScalarExpr>>,
797    ) {
798        let new_inputs = inputs
799            .drain(..)
800            .zip(push_downs)
801            .map(|(input, push_down)| {
802                if !push_down.is_empty() {
803                    input.filter(push_down)
804                } else {
805                    input
806                }
807            })
808            .collect();
809        *inputs = new_inputs;
810    }
811
812    // Checks `get_predicates` to see whether we can push a predicate into the Let binding given
813    // by `id` and `value`.
814    // `users` is the list of those expressions from which we will need to remove a predicate that
815    // is being pushed.
816    fn push_into_let_binding(
817        get_predicates: &mut BTreeMap<Id, BTreeSet<MirScalarExpr>>,
818        id: &LocalId,
819        value: &mut MirRelationExpr,
820        users: &mut [&mut MirRelationExpr],
821    ) {
822        if let Some(list) = get_predicates.remove(&Id::Local(*id)) {
823            if !list.is_empty() {
824                // Remove the predicates in `list` from the users.
825                for user in users {
826                    user.visit_pre_mut(|e| {
827                        if let MirRelationExpr::Filter { input, predicates } = e {
828                            if let MirRelationExpr::Get { id: get_id, .. } = **input {
829                                if get_id == Id::Local(*id) {
830                                    predicates.retain(|p| !list.contains(p));
831                                }
832                            }
833                        }
834                    });
835                }
836                // Apply the predicates in `list` to value. Canonicalize
837                // `list` so that plans are always deterministic.
838                let mut list = list.into_iter().collect::<Vec<_>>();
839                mz_expr::canonicalize::canonicalize_predicates(
840                    &mut list,
841                    &value.typ().column_types,
842                );
843                *value = value.take_dangerous().filter(list);
844            }
845        }
846    }
847
848    /// Returns `(<predicates to retain>, <predicates to push at each input>)`.
849    pub fn push_filters_through_join(
850        input_mapper: &JoinInputMapper,
851        equivalences: &Vec<Vec<MirScalarExpr>>,
852        mut predicates: Vec<MirScalarExpr>,
853    ) -> (Vec<MirScalarExpr>, Vec<Vec<MirScalarExpr>>) {
854        let mut push_downs = vec![Vec::new(); input_mapper.total_inputs()];
855        let mut retain = Vec::new();
856
857        for predicate in predicates.drain(..) {
858            // Track if the predicate has been pushed to at least one input.
859            let mut pushed = false;
860            // For each input, try and see if the join
861            // equivalences allow the predicate to be rewritten
862            // in terms of only columns from that input.
863            for (index, push_down) in push_downs.iter_mut().enumerate() {
864                #[allow(deprecated)] // TODO: use `might_error` if possible.
865                if predicate.is_literal_err() || predicate.contains_error_if_null() {
866                    // Do nothing. We don't push down literal errors,
867                    // as we can't know the join will be non-empty.
868                    //
869                    // We also don't want to push anything that involves `error_if_null`. This is
870                    // for the same reason why in theory we shouldn't really push anything that can
871                    // error, assuming that we want to preserve error semantics. (Because we would
872                    // create a spurious error if some other Join input ends up empty.) We can't fix
873                    // this problem in general (as we can't just not push anything that might
874                    // error), but we decided to fix the specific problem instance involving
875                    // `error_if_null`, because it was very painful:
876                    // <https://github.com/MaterializeInc/database-issues/issues/6258>
877                } else {
878                    let mut localized = predicate.clone();
879                    if input_mapper.try_localize_to_input_with_bound_expr(
880                        &mut localized,
881                        index,
882                        equivalences,
883                    ) {
884                        push_down.push(localized);
885                        pushed = true;
886                    } else if let Some(consequence) = input_mapper
887                        // (`consequence_for_input` assumes that
888                        // `try_localize_to_input_with_bound_expr` has already
889                        // been called on `localized`.)
890                        .consequence_for_input(&localized, index)
891                    {
892                        push_down.push(consequence);
893                        // We don't set `pushed` here! We want to retain the
894                        // predicate, because we only pushed a consequence of
895                        // it, but not the full predicate.
896                    }
897                }
898            }
899
900            if !pushed {
901                retain.push(predicate);
902            }
903        }
904
905        (retain, push_downs)
906    }
907
908    /// Computes "safe" predicates to push through a Map.
909    ///
910    /// In the case of a Filter { Map {...} }, we can always push down the Filter
911    /// by inlining expressions from the Map. We don't want to do this in general,
912    /// however, since general inlining can result in exponential blowup in the size
913    /// of expressions, so we only do this in the case where the size after inlining
914    /// is below a certain limit.
915    ///
916    /// Returns the predicates that can be pushed down, followed by ones that cannot.
917    pub fn push_filters_through_map(
918        map_exprs: &Vec<MirScalarExpr>,
919        predicates: &mut Vec<MirScalarExpr>,
920        input_arity: usize,
921        all_errors: bool,
922    ) -> Result<(Vec<MirScalarExpr>, Vec<MirScalarExpr>), TransformError> {
923        let mut pushdown = Vec::new();
924        let mut retained = Vec::new();
925        for predicate in predicates.drain(..) {
926            // We don't push down literal errors, unless all predicates are.
927            if !predicate.is_literal_err() || all_errors {
928                // Consider inlining Map expressions.
929                if let Some(cleaned) =
930                    Self::inline_if_not_too_big(&predicate, input_arity, map_exprs)?
931                {
932                    pushdown.push(cleaned);
933                } else {
934                    retained.push(predicate);
935                }
936            } else {
937                retained.push(predicate);
938            }
939        }
940        Ok((retained, pushdown))
941    }
942
943    /// This fn should be called with a Filter `expr` that is after a Map. `input_arity` is the
944    /// arity of the input of the Map. This fn eliminates such column refs in `expr` that refer not
945    /// to a column in the input of the Map, but to a column that is created by the Map. It does
946    /// this by transitively inlining Map expressions until no such expression remains that points
947    /// to a Map expression. The return value is the cleaned up expression. The fn bails out with a
948    /// None if the resulting expression would be made too big by the inlinings.
949    ///
950    /// OOO: (Optimizer Optimization Opportunity) This function might do work proportional to the
951    /// total size of the Map expressions. We call this function for each predicate above the Map,
952    /// which will be kind of quadratic, i.e., if there are many predicates and a big Map, then this
953    /// will be slow. We could instead pass a vector of Map expressions and call this fn only once.
954    /// The only downside would be that then the inlining limit being hit in the middle part of this
955    /// function would prevent us from inlining any predicates, even ones that wouldn't hit the
956    /// inlining limit if considered on their own.
957    fn inline_if_not_too_big(
958        expr: &MirScalarExpr,
959        input_arity: usize,
960        map_exprs: &Vec<MirScalarExpr>,
961    ) -> Result<Option<MirScalarExpr>, RecursionLimitError> {
962        let size_limit = 1000;
963
964        // Transitively determine the support of `expr` produced by `map_exprs`
965        // that needs to be inlined.
966        let cols_to_inline = {
967            let mut support = BTreeSet::new();
968
969            // Seed with `map_exprs` support in `expr`.
970            expr.visit_pre(|e| {
971                if let MirScalarExpr::Column(c, _) = e {
972                    if *c >= input_arity {
973                        support.insert(*c);
974                    }
975                }
976            });
977
978            // Compute transitive closure of supports in `map_exprs`.
979            let mut workset = support.iter().cloned().collect::<Vec<_>>();
980            let mut buffer = vec![];
981            while !workset.is_empty() {
982                // Swap the (empty) `drained` buffer with the `workset`.
983                std::mem::swap(&mut workset, &mut buffer);
984                // Drain the `buffer` and update `support` and `workset`.
985                for c in buffer.drain(..) {
986                    map_exprs[c - input_arity].visit_pre(|e| {
987                        if let MirScalarExpr::Column(c, _) = e {
988                            if *c >= input_arity {
989                                if support.insert(*c) {
990                                    workset.push(*c);
991                                }
992                            }
993                        }
994                    });
995                }
996            }
997            support
998        };
999
1000        let mut inlined = BTreeMap::<usize, (MirScalarExpr, usize)>::new();
1001        // Populate the memo table in ascending column order (which respects the
1002        // dependency order of `map_exprs` references). Break early if memoization
1003        // fails for one of the columns in `cols_to_inline`.
1004        for c in cols_to_inline.iter() {
1005            let mut new_expr = map_exprs[*c - input_arity].clone();
1006            let mut new_size = 0;
1007            new_expr.visit_mut_post(&mut |expr| {
1008                new_size += 1;
1009                if let MirScalarExpr::Column(c, _) = expr {
1010                    if *c >= input_arity && new_size <= size_limit {
1011                        // (inlined[c] is safe, because we proceed in column order, and we break out
1012                        // of the loop when we stop inserting into memo.)
1013                        let (m_expr, m_size): &(MirScalarExpr, _) = &inlined[c];
1014                        *expr = m_expr.clone();
1015                        new_size += m_size - 1; // Adjust for the +1 above.
1016                    }
1017                }
1018            })?;
1019
1020            if new_size <= size_limit {
1021                inlined.insert(*c, (new_expr, new_size));
1022            } else {
1023                break;
1024            }
1025        }
1026
1027        // Try to resolve expr against the memo table.
1028        if inlined.len() < cols_to_inline.len() {
1029            Ok(None) // We couldn't memoize all map expressions within the given limit.
1030        } else {
1031            let mut new_expr = expr.clone();
1032            let mut new_size = 0;
1033            new_expr.visit_mut_post(&mut |expr| {
1034                new_size += 1;
1035                if let MirScalarExpr::Column(c, _) = expr {
1036                    if *c >= input_arity && new_size <= size_limit {
1037                        // (inlined[c] is safe because of the outer if condition.)
1038                        let (m_expr, m_size): &(MirScalarExpr, _) = &inlined[c];
1039                        *expr = m_expr.clone();
1040                        new_size += m_size - 1; // Adjust for the +1 above.
1041                    }
1042                }
1043            })?;
1044
1045            soft_assert_eq_no_log!(new_size, new_expr.size());
1046            if new_size <= size_limit {
1047                Ok(Some(new_expr)) // We managed to stay within the limit.
1048            } else {
1049                Ok(None) // Limit exceeded.
1050            }
1051        }
1052    }
1053    // fn inline_if_not_too_big(
1054    //     expr: &MirScalarExpr,
1055    //     input_arity: usize,
1056    //     map_exprs: &Vec<MirScalarExpr>,
1057    // ) -> Result<Option<MirScalarExpr>, RecursionLimitError> {
1058    //     let size_limit = 1000;
1059    //     // Memoize cleaned up versions of Map expressions. (Not necessarily all the Map expressions
1060    //     // will be involved.)
1061    //     let mut memo: BTreeMap<MirScalarExpr, MirScalarExpr> = BTreeMap::new();
1062    //     fn rec(
1063    //         expr: &MirScalarExpr,
1064    //         input_arity: usize,
1065    //         map_exprs: &Vec<MirScalarExpr>,
1066    //         memo: &mut BTreeMap<MirScalarExpr, MirScalarExpr>,
1067    //         size_limit: usize,
1068    //     ) -> Result<Option<MirScalarExpr>, RecursionLimitError> {
1069    //         // (We can't use Entry::or_insert_with, because the closure would need to be fallible.
1070    //         // We also can't manually match on the result of memo.entry, because that holds a
1071    //         // borrow of memo, but we need to pass memo to the recursive call in the middle.)
1072    //         match memo.get(expr) {
1073    //             Some(memoized_result) => Ok(Some(memoized_result.clone())),
1074    //             None => {
1075    //                 let mut expr_size = expr.size()?;
1076    //                 let mut cleaned_expr = expr.clone();
1077    //                 let mut bail = false;
1078    //                 cleaned_expr.try_visit_mut_post(&mut |expr| {
1079    //                     Ok(if !bail {
1080    //                         match expr {
1081    //                             MirScalarExpr::Column(col) => {
1082    //                                 if *col >= input_arity {
1083    //                                     let to_inline = rec(
1084    //                                         &map_exprs[*col - input_arity],
1085    //                                         input_arity,
1086    //                                         map_exprs,
1087    //                                         memo,
1088    //                                         size_limit,
1089    //                                     )?;
1090    //                                     if let Some(to_inline) = to_inline {
1091    //                                         // The `-1` is because the expression that we are
1092    //                                         // replacing has a size of 1.
1093    //                                         expr_size += to_inline.size()? - 1;
1094    //                                         *expr = to_inline;
1095    //                                         if expr_size > size_limit {
1096    //                                             bail = true;
1097    //                                         }
1098    //                                     } else {
1099    //                                         bail = true;
1100    //                                     }
1101    //                                 }
1102    //                             }
1103    //                             _ => (),
1104    //                         }
1105    //                     })
1106    //                 })?;
1107    //                 soft_assert_eq!(cleaned_expr.size()?, expr_size);
1108    //                 if !bail {
1109    //                     memo.insert(expr.clone(), cleaned_expr.clone());
1110    //                     Ok(Some(cleaned_expr))
1111    //                 } else {
1112    //                     Ok(None)
1113    //                 }
1114    //             }
1115    //         }
1116    //     }
1117    //     rec(expr, input_arity, map_exprs, &mut memo, size_limit)
1118    // }
1119
1120    /// Computes "safe" predicate to push through a FlatMap.
1121    ///
1122    /// In the case of a Filter { FlatMap {...} }, we want to push through all predicates
1123    /// that (1) are not literal errors and (2) have support exclusively in the columns
1124    /// provided by the FlatMap input.
1125    ///
1126    /// Returns the predicates that can be pushed down, followed by ones that cannot.
1127    fn push_filters_through_flat_map(
1128        predicates: &mut Vec<MirScalarExpr>,
1129        input_arity: usize,
1130    ) -> (Vec<MirScalarExpr>, Vec<MirScalarExpr>) {
1131        let mut pushdown = Vec::new();
1132        let mut retained = Vec::new();
1133        for predicate in predicates.drain(..) {
1134            // First, check if we can push this predicate down. We can do so if and only if:
1135            // (1) the predicate is not a literal error, and
1136            // (2) each column it references is from the input.
1137            if (!predicate.is_literal_err()) && predicate.support().iter().all(|c| *c < input_arity)
1138            {
1139                pushdown.push(predicate);
1140            } else {
1141                retained.push(predicate);
1142            }
1143        }
1144        (retained, pushdown)
1145    }
1146
1147    /// If `s` is of the form
1148    /// `(isnull(expr1) && isnull(expr2)) || (expr1 = expr2)`, or
1149    /// `(decompose_is_null(expr1) && decompose_is_null(expr2)) || (expr1 = expr2)`,
1150    /// extract `expr1` and `expr2`.
1151    fn extract_equal_or_both_null(
1152        s: &mut MirScalarExpr,
1153        column_types: &[SqlColumnType],
1154    ) -> Option<(MirScalarExpr, MirScalarExpr)> {
1155        if let MirScalarExpr::CallVariadic {
1156            func: VariadicFunc::Or,
1157            exprs,
1158        } = s
1159        {
1160            if let &[ref or_lhs, ref or_rhs] = &**exprs {
1161                // Check both orders of operands of the OR
1162                return Self::extract_equal_or_both_null_inner(or_lhs, or_rhs, column_types)
1163                    .or_else(|| {
1164                        Self::extract_equal_or_both_null_inner(or_rhs, or_lhs, column_types)
1165                    });
1166            }
1167        }
1168        None
1169    }
1170
1171    fn extract_equal_or_both_null_inner(
1172        or_arg1: &MirScalarExpr,
1173        or_arg2: &MirScalarExpr,
1174        column_types: &[SqlColumnType],
1175    ) -> Option<(MirScalarExpr, MirScalarExpr)> {
1176        use mz_expr::BinaryFunc;
1177        if let MirScalarExpr::CallBinary {
1178            func: BinaryFunc::Eq,
1179            expr1: eq_lhs,
1180            expr2: eq_rhs,
1181        } = &or_arg2
1182        {
1183            let isnull1 = eq_lhs.clone().call_is_null();
1184            let isnull2 = eq_rhs.clone().call_is_null();
1185            let both_null = MirScalarExpr::CallVariadic {
1186                func: VariadicFunc::And,
1187                exprs: vec![isnull1, isnull2],
1188            };
1189
1190            if Self::extract_reduced_conjunction_terms(both_null, column_types)
1191                == Self::extract_reduced_conjunction_terms(or_arg1.clone(), column_types)
1192            {
1193                return Some(((**eq_lhs).clone(), (**eq_rhs).clone()));
1194            }
1195        }
1196        None
1197    }
1198
1199    /// Reduces the given expression and returns its AND-ed terms.
1200    fn extract_reduced_conjunction_terms(
1201        mut s: MirScalarExpr,
1202        column_types: &[SqlColumnType],
1203    ) -> Vec<MirScalarExpr> {
1204        s.reduce(column_types);
1205
1206        if let MirScalarExpr::CallVariadic {
1207            func: VariadicFunc::And,
1208            exprs,
1209        } = s
1210        {
1211            exprs
1212        } else {
1213            vec![s]
1214        }
1215    }
1216}