mz_transform/
literal_constraints.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//! See if there are predicates of the form `<expr> = literal` that can be sped up using an index.
11//! More specifically, look for an MFP on top of a Get, where the MFP has an appropriate filter, and
12//! the Get has a matching index. Convert these to `IndexedFilter` joins, which is a semi-join with
13//! a constant collection.
14//!
15//! E.g.: Logically, we go from something like
16//! `SELECT f1, f2, f3 FROM t WHERE t.f1 = lit1 AND t.f2 = lit2`
17//! to
18//! `SELECT f1, f2, f3 FROM t, (SELECT * FROM (VALUES (lit1, lit2))) as filter_list
19//!  WHERE t.f1 = filter_list.column1 AND t.f2 = filter_list.column2`
20
21use std::collections::{BTreeMap, BTreeSet};
22
23use itertools::Itertools;
24use mz_expr::JoinImplementation::IndexedFilter;
25use mz_expr::canonicalize::canonicalize_predicates;
26use mz_expr::visit::{Visit, VisitChildren};
27use mz_expr::{BinaryFunc, Id, MapFilterProject, MirRelationExpr, MirScalarExpr, VariadicFunc};
28use mz_ore::collections::CollectionExt;
29use mz_ore::iter::IteratorExt;
30use mz_ore::stack::RecursionLimitError;
31use mz_ore::vec::swap_remove_multiple;
32use mz_repr::{Diff, GlobalId, RelationType, Row};
33
34use crate::TransformCtx;
35use crate::canonicalize_mfp::CanonicalizeMfp;
36use crate::notice::IndexTooWideForLiteralConstraints;
37
38/// Convert literal constraints into `IndexedFilter` joins.
39#[derive(Debug)]
40pub struct LiteralConstraints;
41
42impl crate::Transform for LiteralConstraints {
43    fn name(&self) -> &'static str {
44        "LiteralConstraints"
45    }
46
47    #[mz_ore::instrument(
48        target = "optimizer",
49        level = "debug",
50        fields(path.segment = "literal_constraints")
51    )]
52    fn actually_perform_transform(
53        &self,
54        relation: &mut MirRelationExpr,
55        ctx: &mut TransformCtx,
56    ) -> Result<(), crate::TransformError> {
57        let result = self.action(relation, ctx);
58        mz_repr::explain::trace_plan(&*relation);
59        result
60    }
61}
62
63impl LiteralConstraints {
64    fn action(
65        &self,
66        relation: &mut MirRelationExpr,
67        transform_ctx: &mut TransformCtx,
68    ) -> Result<(), crate::TransformError> {
69        let mut mfp = MapFilterProject::extract_non_errors_from_expr_mut(relation);
70        relation.try_visit_mut_children(|e| self.action(e, transform_ctx))?;
71
72        if let MirRelationExpr::Get {
73            id: Id::Global(id),
74            ref typ,
75            ..
76        } = *relation
77        {
78            let orig_mfp = mfp.clone();
79
80            // Preparation for the literal constraints detection.
81            Self::inline_literal_constraints(&mut mfp);
82            Self::list_of_predicates_to_and_of_predicates(&mut mfp);
83            Self::distribute_and_over_or(&mut mfp)?;
84            Self::unary_and(&mut mfp);
85
86            /// The above preparation might make the MFP more complicated, so we'll later want to
87            /// either undo the preparation transformations or get back to `orig_mfp`.
88            fn undo_preparation(
89                mfp: &mut MapFilterProject,
90                orig_mfp: &MapFilterProject,
91                relation: &MirRelationExpr,
92                relation_type: RelationType,
93            ) {
94                // undo list_of_predicates_to_and_of_predicates, distribute_and_over_or, unary_and
95                // (It undoes the latter 2 through `MirScalarExp::reduce`.)
96                LiteralConstraints::canonicalize_predicates(mfp, relation, relation_type);
97                // undo inline_literal_constraints
98                mfp.optimize();
99                // We can usually undo, but sometimes not (see comment on `distribute_and_over_or`),
100                // so in those cases we might have a more complicated MFP than the original MFP
101                // (despite the removal of the literal constraints and/or contradicting OR args).
102                // So let's use the simpler one.
103                if LiteralConstraints::predicates_size(orig_mfp)
104                    < LiteralConstraints::predicates_size(mfp)
105                {
106                    *mfp = orig_mfp.clone();
107                }
108            }
109
110            let removed_contradicting_or_args = Self::remove_impossible_or_args(&mut mfp)?;
111
112            // todo: We might want to also call `canonicalize_equivalences`,
113            // see near the end of literal_constraints.slt.
114
115            let inp_typ = typ.clone();
116
117            let key_val = Self::detect_literal_constraints(&mfp, id, transform_ctx);
118
119            match key_val {
120                None => {
121                    // We didn't find a usable index, so no chance to remove literal constraints.
122                    // But, we might have removed contradicting OR args.
123                    if removed_contradicting_or_args {
124                        undo_preparation(&mut mfp, &orig_mfp, relation, inp_typ);
125                    } else {
126                        // We didn't remove anything, so let's go with the original MFP.
127                        mfp = orig_mfp;
128                    }
129                }
130                Some((idx_id, key, possible_vals)) => {
131                    // We found a usable index. We'll try to remove the corresponding literal
132                    // constraints.
133                    if Self::remove_literal_constraints(&mut mfp, &key)
134                        || removed_contradicting_or_args
135                    {
136                        // We were able to remove the literal constraints or contradicting OR args,
137                        // so we would like to use this new MFP, so we try undoing the preparation.
138                        undo_preparation(&mut mfp, &orig_mfp, relation, inp_typ.clone());
139                    } else {
140                        // We were not able to remove the literal constraint, so `mfp` is
141                        // equivalent to `orig_mfp`, but `orig_mfp` is often simpler (or the same).
142                        mfp = orig_mfp;
143                    }
144
145                    // We transform the Get into a semi-join with a constant collection.
146
147                    let inp_id = id.clone();
148                    let filter_list = MirRelationExpr::Constant {
149                        rows: Ok(possible_vals
150                            .iter()
151                            .map(|val| (val.clone(), Diff::ONE))
152                            .collect()),
153                        typ: mz_repr::RelationType {
154                            column_types: key
155                                .iter()
156                                .map(|e| {
157                                    e.typ(&inp_typ.column_types)
158                                        // We make sure to not include a null in `expr_eq_literal`.
159                                        .nullable(false)
160                                })
161                                .collect(),
162                            // (Note that the key inference for `MirRelationExpr::Constant` inspects
163                            // the constant values to detect keys not listed within the node, but it
164                            // can only detect a single-column key this way. A multi-column key is
165                            // common here, so we explicitly add it.)
166                            keys: vec![(0..key.len()).collect()],
167                        },
168                    }
169                    .arrange_by(&[(0..key.len()).map(MirScalarExpr::Column).collect_vec()]);
170
171                    if possible_vals.is_empty() {
172                        // Even better than what we were hoping for: Found contradicting
173                        // literal constraints, so the whole relation is empty.
174                        relation.take_safely(Some(inp_typ));
175                    } else {
176                        // The common case: We need to build the join which is the main point of
177                        // this transform.
178                        *relation = MirRelationExpr::Join {
179                            // It's important to keep the `filter_list` in the second position.
180                            // Both the lowering and EXPLAIN depends on this.
181                            inputs: vec![relation.clone().arrange_by(&[key.clone()]), filter_list],
182                            equivalences: key
183                                .iter()
184                                .enumerate()
185                                .map(|(i, e)| {
186                                    vec![(*e).clone(), MirScalarExpr::column(i + inp_typ.arity())]
187                                })
188                                .collect(),
189                            implementation: IndexedFilter(
190                                inp_id,
191                                idx_id,
192                                key.clone(),
193                                possible_vals,
194                            ),
195                        };
196
197                        // Rebuild the MFP to add the projection that removes the columns coming from
198                        // the filter_list side of the join.
199                        let (map, filter, project) = mfp.as_map_filter_project();
200                        mfp = MapFilterProject::new(inp_typ.arity() + key.len())
201                            .project(0..inp_typ.arity()) // make the join semi
202                            .map(map)
203                            .filter(filter)
204                            .project(project);
205                        mfp.optimize()
206                    }
207                }
208            }
209        }
210
211        CanonicalizeMfp::rebuild_mfp(mfp, relation);
212
213        Ok(())
214    }
215
216    /// Detects literal constraints in an MFP on top of a Get of `id`, and a matching index that can
217    /// be used to speed up the Filter of the MFP.
218    ///
219    /// For example, if there is an index on `(f1, f2)`, and the Filter is
220    /// `(f1 = 3 AND f2 = 5) OR (f1 = 7 AND f2 = 9)`, it returns `Some([f1, f2], [[3,5], [7,9]])`.
221    ///
222    /// We can use an index if each argument of the OR includes a literal constraint on each of the
223    /// key fields of the index. Extra predicates inside the OR arguments are ok.
224    ///
225    /// Returns (idx_id, idx_key, values to lookup in the index).
226    fn detect_literal_constraints(
227        mfp: &MapFilterProject,
228        get_id: GlobalId,
229        transform_ctx: &mut TransformCtx,
230    ) -> Option<(GlobalId, Vec<MirScalarExpr>, Vec<Row>)> {
231        // Checks whether an index with the specified key can be used to speed up the given filter.
232        // See comment of `IndexMatch`.
233        fn match_index(key: &[MirScalarExpr], or_args: &Vec<MirScalarExpr>) -> IndexMatch {
234            if key.is_empty() {
235                // Nothing to do with an index that has an empty key.
236                return IndexMatch::UnusableNoSubset;
237            }
238            if !key.iter().all_unique() {
239                // This is a weird index. Why does it have duplicate key expressions?
240                return IndexMatch::UnusableNoSubset;
241            }
242            let mut literal_values = Vec::new();
243            let mut inv_cast_any = false;
244            // This starts with all key fields of the index.
245            // At the end, it will contain a subset S of index key fields such that if the index had
246            // only S as its key, then the index would be usable.
247            let mut usable_key_fields = key.iter().collect::<BTreeSet<_>>();
248            let mut usable = true;
249            for or_arg in or_args {
250                let mut row = Row::default();
251                let mut packer = row.packer();
252                for key_field in key {
253                    let and_args = or_arg.and_or_args(VariadicFunc::And);
254                    // Let's find a constraint for this key field
255                    if let Some((literal, inv_cast)) = and_args
256                        .iter()
257                        .find_map(|and_arg| and_arg.expr_eq_literal(key_field))
258                    {
259                        // (Note that the above find_map can find only 0 or 1 result, because
260                        // of `remove_impossible_or_args`.)
261                        packer.push(literal.unpack_first());
262                        inv_cast_any |= inv_cast;
263                    } else {
264                        // There is an `or_arg` where we didn't find a constraint for a key field,
265                        // so the index is unusable. Throw out the field from the usable fields.
266                        usable = false;
267                        usable_key_fields.remove(key_field);
268                        if usable_key_fields.is_empty() {
269                            return IndexMatch::UnusableNoSubset;
270                        }
271                    }
272                }
273                literal_values.push(row);
274            }
275            if usable {
276                // We should deduplicate, because a constraint can be duplicated by
277                // `distribute_and_over_or`. For example: `IN ('l1', 'l2') AND (a > 0 OR a < 5)`:
278                // the 2 args of the OR will cause the IN constraints to be duplicated. This doesn't
279                // alter the meaning of the expression when evaluated as a filter, but if we extract
280                // those literals 2 times into `literal_values` then the Peek code will look up
281                // those keys from the index 2 times, leading to duplicate results.
282                literal_values.sort();
283                literal_values.dedup();
284                IndexMatch::Usable(literal_values, inv_cast_any)
285            } else {
286                if usable_key_fields.is_empty() {
287                    IndexMatch::UnusableNoSubset
288                } else {
289                    IndexMatch::UnusableTooWide(
290                        usable_key_fields.into_iter().cloned().collect_vec(),
291                    )
292                }
293            }
294        }
295
296        let or_args = Self::get_or_args(mfp);
297
298        let index_matches = transform_ctx
299            .indexes
300            .indexes_on(get_id)
301            .map(|(index_id, key)| (index_id, key.to_owned(), match_index(key, &or_args)))
302            .collect_vec();
303
304        let result = index_matches
305            .iter()
306            .cloned()
307            .filter_map(|(idx_id, key, index_match)| match index_match {
308                IndexMatch::Usable(vals, inv_cast) => Some((idx_id, key, vals, inv_cast)),
309                _ => None,
310            })
311            // Maximize:
312            //  1. number of predicates that are sped using a single index.
313            //  2. whether we are using a simpler index by having removed a cast from the key expr.
314            .max_by_key(|(_idx_id, key, _vals, inv_cast)| (key.len(), *inv_cast))
315            .map(|(idx_id, key, vals, _inv_cast)| (idx_id, key, vals));
316
317        if result.is_none() && !or_args.is_empty() {
318            // Let's see if we can give a hint to the user.
319            index_matches
320                .into_iter()
321                .for_each(|(index_id, index_key, index_match)| {
322                    match index_match {
323                        IndexMatch::UnusableTooWide(usable_subset) => {
324                            // see comment of `UnusableTooWide`
325                            assert!(!usable_subset.is_empty());
326                            // Determine literal values that we would get if the index was on
327                            // `usable_subset`.
328                            let literal_values = match match_index(&usable_subset, &or_args) {
329                                IndexMatch::Usable(literal_vals, _) => literal_vals,
330                                _ => unreachable!(), // `usable_subset` would make the index usable.
331                            };
332
333                            // Let's come up with a recommendation for what columns to index:
334                            // Intersect literal constraints across all OR args. (Which might
335                            // include columns that are NOT in this index, and therefore not in
336                            // `usable_subset`.)
337                            let recommended_key = or_args
338                                .iter()
339                                .map(|or_arg| {
340                                    let and_args = or_arg.and_or_args(VariadicFunc::And);
341                                    and_args
342                                        .iter()
343                                        .filter_map(|and_arg| and_arg.any_expr_eq_literal())
344                                        .collect::<BTreeSet<_>>()
345                                })
346                                .reduce(|fields1, fields2| {
347                                    fields1.intersection(&fields2).cloned().collect()
348                                })
349                                // The unwrap is safe because above we checked `!or_args.is_empty()`
350                                .unwrap()
351                                .into_iter()
352                                .collect_vec();
353
354                            transform_ctx.df_meta.push_optimizer_notice_dedup(
355                                IndexTooWideForLiteralConstraints {
356                                    index_id,
357                                    index_key,
358                                    usable_subset,
359                                    literal_values,
360                                    index_on_id: get_id,
361                                    recommended_key,
362                                },
363                            )
364                        }
365                        _ => (),
366                    }
367                });
368        }
369
370        result
371    }
372
373    /// Removes the expressions that [LiteralConstraints::detect_literal_constraints] found, if
374    /// possible. Returns whether it removed anything.
375    /// For example, if the key of the detected literal constraint is just `f1`, and we have the
376    /// expression
377    /// `(f1 = 3 AND f2 = 5) OR (f1 = 7 AND f2 = 5)`, then this modifies it to `f2 = 5`.
378    /// However, if OR branches differ in their non-key parts, then we cannot remove the literal
379    /// constraint. For example,
380    /// `(f1 = 3 AND f2 = 5) OR (f1 = 7 AND f2 = 555)`, then we cannot remove the `f1` parts,
381    /// because then the filter wouldn't know whether to check `f2 = 5` or `f2 = 555`.
382    fn remove_literal_constraints(mfp: &mut MapFilterProject, key: &Vec<MirScalarExpr>) -> bool {
383        let or_args = Self::get_or_args(mfp);
384        if or_args.len() == 0 {
385            return false;
386        }
387
388        // In simple situations it would be enough to check here that if we remove the detected
389        // literal constraints from each OR arg, then the residual OR args are all equal.
390        // However, this wouldn't be able to perform the removal when the expression that should
391        // remain in the end has an OR. This is because conversion to DNF makes duplicates of
392        // every literal constraint, with different residuals. To also handle this case, we collect
393        // the possible residuals for every literal constraint row, and check that all sets are
394        // equal. Example: The user wrote
395        // `WHERE ((a=1 AND b=1) OR (a=2 AND b=2)) AND (c OR (d AND e))`.
396        // The DNF of this is
397        // `(a=1 AND b=1 AND c) OR (a=1 AND b=1 AND d AND e) OR (a=2 AND b=2 AND c) OR (a=2 AND b=2 AND d AND e)`.
398        // Then `constraints_to_residual_sets` will be:
399        // [
400        //   [`a=1`, `b=1`]  ->  {[`c`], [`d`, `e`]},
401        //   [`a=2`, `b=2`]  ->  {[`c`], [`d`, `e`]}
402        // ]
403        // After removing the literal constraints we have
404        // `c OR (d AND e)`
405        let mut constraints_to_residual_sets = BTreeMap::new();
406        or_args.iter().for_each(|or_arg| {
407            let and_args = or_arg.and_or_args(VariadicFunc::And);
408            let (mut constraints, mut residual): (Vec<_>, Vec<_>) =
409                and_args.iter().cloned().partition(|and_arg| {
410                    key.iter()
411                        .any(|key_field| matches!(and_arg.expr_eq_literal(key_field), Some(..)))
412                });
413            // In every or_arg there has to be some literal constraints, otherwise
414            // `detect_literal_constraints` would have returned None.
415            assert!(constraints.len() >= 1);
416            // `remove_impossible_or_args` made sure that inside each or_arg, each
417            // expression can be literal constrained only once. So if we find one of the
418            // key fields being literal constrained, then it's definitely that literal
419            // constraint that detect_literal_constraints based one of its return values on.
420            //
421            // This is important, because without `remove_impossible_or_args`, we might
422            // have the situation here that or_arg would be something like
423            // `a = 5 AND a = 8`, of which `detect_literal_constraints` found only the `a = 5`,
424            // but here we would remove both the `a = 5` and the `a = 8`.
425            constraints.sort();
426            residual.sort();
427            let entry = constraints_to_residual_sets
428                .entry(constraints)
429                .or_insert_with(BTreeSet::new);
430            entry.insert(residual);
431        });
432        let residual_sets = constraints_to_residual_sets
433            .into_iter()
434            .map(|(_constraints, residual_set)| residual_set)
435            .collect::<Vec<_>>();
436        if residual_sets.iter().all_equal() {
437            // We can remove the literal constraint
438            assert!(residual_sets.len() >= 1); // We already checked `or_args.len() == 0` above
439            let residual_set = residual_sets.into_iter().into_first();
440            let new_pred = MirScalarExpr::CallVariadic {
441                func: VariadicFunc::Or,
442                exprs: residual_set
443                    .into_iter()
444                    .map(|residual| MirScalarExpr::CallVariadic {
445                        func: VariadicFunc::And,
446                        exprs: residual,
447                    })
448                    .collect::<Vec<_>>(),
449            };
450            let (map, _predicates, project) = mfp.as_map_filter_project();
451            *mfp = MapFilterProject::new(mfp.input_arity)
452                .map(map)
453                .filter(std::iter::once(new_pred))
454                .project(project);
455
456            true
457        } else {
458            false
459        }
460    }
461
462    /// 1. Removes such OR args in which there are contradicting literal constraints.
463    /// 2. Also, if an OR arg doesn't have any contradiction, this fn just deduplicates
464    /// the AND arg list of that OR arg. (Might additionally sort all AND arg lists.)
465    ///
466    /// Returns whether it performed any removal or deduplication.
467    ///
468    /// Example for 1:
469    /// `<arg1> OR (a = 5 AND a = 5 AND a = 8) OR <arg3>`
470    /// -->
471    /// `<arg1> OR <arg3> `
472    ///
473    /// Example for 2:
474    /// `<arg1> OR (a = 5 AND a = 5 AND b = 8) OR <arg3>`
475    /// -->
476    /// `<arg1> OR (a = 5 AND b = 8) OR <arg3>`
477    fn remove_impossible_or_args(mfp: &mut MapFilterProject) -> Result<bool, RecursionLimitError> {
478        let mut or_args = Self::get_or_args(mfp);
479        if or_args.len() == 0 {
480            return Ok(false);
481        }
482        let mut to_remove = Vec::new();
483        let mut changed = false;
484        or_args.iter_mut().enumerate().for_each(|(i, or_arg)| {
485            if let MirScalarExpr::CallVariadic {
486                func: VariadicFunc::And,
487                exprs: and_args,
488            } = or_arg
489            {
490                if and_args
491                    .iter()
492                    .any(|e| e.impossible_literal_equality_because_types())
493                {
494                    changed = true;
495                    to_remove.push(i);
496                } else {
497                    and_args.sort_by_key(|e: &MirScalarExpr| e.invert_casts_on_expr_eq_literal());
498                    let and_args_before_dedup = and_args.clone();
499                    and_args
500                        .dedup_by_key(|e: &mut MirScalarExpr| e.invert_casts_on_expr_eq_literal());
501                    if *and_args != and_args_before_dedup {
502                        changed = true;
503                    }
504                    // Deduplicated, so we cannot have something like `a = 5 AND a = 5`.
505                    // This means that if we now have `<expr1> = <literal1> AND <expr1> = <literal2>`,
506                    // then `literal1` is definitely not the same as `literal2`. This means that this
507                    // whole or_arg is a contradiction, because it's something like `a = 5 AND a = 8`.
508                    let mut literal_constrained_exprs = and_args
509                        .iter()
510                        .filter_map(|and_arg| and_arg.any_expr_eq_literal());
511                    if !literal_constrained_exprs.all_unique() {
512                        changed = true;
513                        to_remove.push(i);
514                    }
515                }
516            } else {
517                // `unary_and` made sure that each OR arg is an AND
518                unreachable!("OR arg was not an AND in remove_impossible_or_args");
519            }
520        });
521        // We remove the marked OR args.
522        // (If the OR has 0 or 1 args remaining, then `reduce_and_canonicalize_and_or` will later
523        // further simplify.)
524        swap_remove_multiple(&mut or_args, to_remove);
525        // Rebuild the MFP if needed
526        if changed {
527            let new_predicates = vec![MirScalarExpr::CallVariadic {
528                func: VariadicFunc::Or,
529                exprs: or_args,
530            }];
531            let (map, _predicates, project) = mfp.as_map_filter_project();
532            *mfp = MapFilterProject::new(mfp.input_arity)
533                .map(map)
534                .filter(new_predicates)
535                .project(project);
536            Ok(true)
537        } else {
538            Ok(false)
539        }
540    }
541
542    /// Returns the arguments of the predicate's top-level OR as a Vec.
543    /// If there is no top-level OR, then interpret the predicate as a 1-arg OR, i.e., return a
544    /// 1-element Vec.
545    ///
546    /// Assumes that [LiteralConstraints::list_of_predicates_to_and_of_predicates] has already run.
547    fn get_or_args(mfp: &MapFilterProject) -> Vec<MirScalarExpr> {
548        assert_eq!(mfp.predicates.len(), 1); // list_of_predicates_to_and_of_predicates ensured this
549        let (_, pred) = mfp.predicates.get(0).unwrap();
550        pred.and_or_args(VariadicFunc::Or)
551    }
552
553    /// Makes the job of [LiteralConstraints::detect_literal_constraints] easier by undoing some CSE to
554    /// reconstruct literal constraints.
555    fn inline_literal_constraints(mfp: &mut MapFilterProject) {
556        let mut should_inline = vec![false; mfp.input_arity + mfp.expressions.len()];
557        // Mark those expressions for inlining that contain a subexpression of the form
558        // `<xxx> = <lit>` or `<lit> = <xxx>`.
559        for (i, e) in mfp.expressions.iter().enumerate() {
560            e.visit_pre(|s| {
561                if s.any_expr_eq_literal().is_some() {
562                    should_inline[i + mfp.input_arity] = true;
563                }
564            });
565        }
566        // Whenever
567        // `<Column(i)> = <lit>` or `<lit> = <Column(i)>`
568        // appears in a predicate, mark the ith expression to be inlined.
569        for (_before, p) in mfp.predicates.iter() {
570            p.visit_pre(|e| {
571                if let MirScalarExpr::CallBinary {
572                    func: BinaryFunc::Eq,
573                    expr1,
574                    expr2,
575                } = e
576                {
577                    if matches!(**expr1, MirScalarExpr::Literal(..)) {
578                        if let MirScalarExpr::Column(col) = **expr2 {
579                            if col >= mfp.input_arity {
580                                should_inline[col] = true;
581                            }
582                        }
583                    }
584                    if matches!(**expr2, MirScalarExpr::Literal(..)) {
585                        if let MirScalarExpr::Column(col) = **expr1 {
586                            if col >= mfp.input_arity {
587                                should_inline[col] = true;
588                            }
589                        }
590                    }
591                }
592            });
593        }
594        // Perform the marked inlinings.
595        mfp.perform_inlining(should_inline);
596    }
597
598    /// MFPs have a Vec of predicates `[p1, p2, ...]`, which logically represents `p1 AND p2 AND ...`.
599    /// This function performs this conversion. Note that it might create a variadic AND with
600    /// 0 or 1 args, so the resulting predicate Vec always has exactly 1 element.
601    fn list_of_predicates_to_and_of_predicates(mfp: &mut MapFilterProject) {
602        // Rebuild the MFP. (Unfortunately, we cannot modify the predicates in place, because MFP
603        // predicates also have a "before" field, which we need to update. (`filter` will recompute
604        // these.)
605        let (map, _predicates, project) = mfp.as_map_filter_project();
606        let new_predicates = vec![MirScalarExpr::CallVariadic {
607            func: VariadicFunc::And,
608            exprs: mfp.predicates.iter().map(|(_, p)| p.clone()).collect(),
609        }];
610        *mfp = MapFilterProject::new(mfp.input_arity)
611            .map(map)
612            .filter(new_predicates)
613            .project(project);
614    }
615
616    /// Call [mz_expr::canonicalize::canonicalize_predicates] on each of the predicates in the MFP.
617    fn canonicalize_predicates(
618        mfp: &mut MapFilterProject,
619        relation: &MirRelationExpr,
620        relation_type: RelationType,
621    ) {
622        let (map, mut predicates, project) = mfp.as_map_filter_project();
623        let typ_after_map = relation
624            .clone()
625            .map(map.clone())
626            .typ_with_input_types(&[relation_type]);
627        canonicalize_predicates(&mut predicates, &typ_after_map.column_types);
628        // Rebuild the MFP with the new predicates.
629        *mfp = MapFilterProject::new(mfp.input_arity)
630            .map(map)
631            .filter(predicates)
632            .project(project);
633    }
634
635    /// Distribute AND over OR + do flatten_and_or until fixed point.
636    /// This effectively converts to disjunctive normal form (DNF) (i.e., an OR of ANDs), because
637    /// [MirScalarExpr::reduce] did Demorgans and double-negation-elimination. So after
638    /// [MirScalarExpr::reduce], we get here a tree of AND/OR nodes. A distribution step lifts an OR
639    /// up the tree by 1 level, and a [MirScalarExpr::flatten_associative] merges two ORs that are at
640    /// adjacent levels, so eventually we'll end up with just one OR that is at the top of the tree,
641    /// with ANDs below it.
642    /// For example:
643    /// (a || b) && (c || d)
644    ///   ->
645    /// ((a || b) && c) || ((a || b) && d)
646    ///   ->
647    /// (a && c) || (b && c) || (a && d) || (b && d)
648    /// (This is a variadic OR with 4 arguments.)
649    ///
650    /// Example:
651    /// User wrote `WHERE (a,b) IN ((1,2), (1,4), (8,5))`,
652    /// from which [MirScalarExpr::undistribute_and_or] made this before us:
653    /// (#0 = 1 AND (#1 = 2 OR #1 = 4)) OR (#0 = 8 AND #1 = 5)
654    /// And now we distribute the first AND over the first OR in 2 steps: First to
655    /// ((#0 = 1 AND #1 = 2) OR (#0 = 1 AND #1 = 4)) OR (#0 = 8 AND #1 = 5)
656    /// then [MirScalarExpr::flatten_associative]:
657    /// (#0 = 1 AND #1 = 2) OR (#0 = 1 AND #1 = 4) OR (#0 = 8 AND #1 = 5)
658    ///
659    /// Note that [MirScalarExpr::undistribute_and_or] is not exactly an inverse to this because
660    /// 1) it can undistribute both AND over OR and OR over AND.
661    /// 2) it cannot always undo the distribution, because an expression might have multiple
662    /// overlapping undistribution opportunities, see comment there.
663    fn distribute_and_over_or(mfp: &mut MapFilterProject) -> Result<(), RecursionLimitError> {
664        mfp.predicates.iter_mut().try_for_each(|(_, p)| {
665            let mut old_p = MirScalarExpr::column(0);
666            while old_p != *p {
667                let size = p.size();
668                // We might make the expression exponentially larger, so we should have some limit.
669                // Below 1000 (e.g., a single IN list of ~300 elements, or 3 IN lists of 4-5
670                // elements each), we are <10 ms for a single IN list, and even less for multiple IN
671                // lists.
672                if size > 1000 {
673                    break;
674                }
675                old_p = p.clone();
676                p.visit_mut_post(&mut |e: &mut MirScalarExpr| {
677                    if let MirScalarExpr::CallVariadic {
678                        func: VariadicFunc::And,
679                        exprs: and_args,
680                    } = e
681                    {
682                        if let Some((i, _)) = and_args.iter().enumerate().find(|(_i, a)| {
683                            matches!(
684                                a,
685                                MirScalarExpr::CallVariadic {
686                                    func: VariadicFunc::Or,
687                                    ..
688                                }
689                            )
690                        }) {
691                            // We found an AND whose ith argument is an OR. We'll distribute the other
692                            // args of the AND over this OR.
693                            let mut or = and_args.swap_remove(i);
694                            let to_distribute = MirScalarExpr::CallVariadic {
695                                func: VariadicFunc::And,
696                                exprs: (*and_args).clone(),
697                            };
698                            if let MirScalarExpr::CallVariadic {
699                                func: VariadicFunc::Or,
700                                exprs: ref mut or_args,
701                            } = or
702                            {
703                                or_args.iter_mut().for_each(|a| {
704                                    *a = a.clone().and(to_distribute.clone());
705                                });
706                            } else {
707                                unreachable!(); // because the `find` found a match already
708                            }
709                            *e = or; // The modified OR will be the new top-level expr.
710                        }
711                    }
712                })?;
713                p.visit_mut_post(&mut |e: &mut MirScalarExpr| {
714                    e.flatten_associative();
715                })?;
716            }
717            Ok(())
718        })
719    }
720
721    /// For each of the arguments of the top-level OR (if no top-level OR, then interpret the whole
722    /// expression as a 1-arg OR, see [LiteralConstraints::get_or_args]), check if it's an AND, and
723    /// if not, then wrap it in a 1-arg AND.
724    fn unary_and(mfp: &mut MapFilterProject) {
725        let mut or_args = Self::get_or_args(mfp);
726        let mut changed = false;
727        or_args.iter_mut().for_each(|or_arg| {
728            if !matches!(
729                or_arg,
730                MirScalarExpr::CallVariadic {
731                    func: VariadicFunc::And,
732                    ..
733                }
734            ) {
735                *or_arg = MirScalarExpr::CallVariadic {
736                    func: VariadicFunc::And,
737                    exprs: vec![or_arg.clone()],
738                };
739                changed = true;
740            }
741        });
742        if changed {
743            let new_predicates = vec![MirScalarExpr::CallVariadic {
744                func: VariadicFunc::Or,
745                exprs: or_args,
746            }];
747            let (map, _predicates, project) = mfp.as_map_filter_project();
748            *mfp = MapFilterProject::new(mfp.input_arity)
749                .map(map)
750                .filter(new_predicates)
751                .project(project);
752        }
753    }
754
755    fn predicates_size(mfp: &MapFilterProject) -> usize {
756        let mut sum = 0;
757        for (_, p) in mfp.predicates.iter() {
758            sum = sum + p.size();
759        }
760        sum
761    }
762}
763
764/// Whether an index is usable to speed up a Filter with literal constraints.
765#[derive(Clone)]
766enum IndexMatch {
767    /// The index is usable, that is, each OR argument constrains each key field.
768    ///
769    /// The `Vec<Row>` has the constraining literal values, where each Row corresponds to one OR
770    /// argument, and each value in the Row corresponds to one key field.
771    ///
772    /// The `bool` indicates whether we needed to inverse cast equalities to match them up with key
773    /// fields. The inverse cast enables index usage when an implicit cast is wrapping a key field.
774    /// E.g., if `a` is smallint, and the user writes `a = 5`, then HIR inserts an implicit cast:
775    /// `smallint_to_integer(a) = 5`, which we invert to `a = 5`, where the `5` is a smallint
776    /// literal. For more details on the inversion, see `invert_casts_on_expr_eq_literal_inner`.
777    Usable(Vec<Row>, bool),
778    /// The index is unusable. However, there is a subset of key fields such that if the index would
779    /// be only on this subset, then it would be usable.
780    /// Note: this Vec is never empty. (If it were empty, then we'd get `UnusableNoSubset` instead.)
781    UnusableTooWide(Vec<MirScalarExpr>),
782    /// The index is unusable. Moreover, none of its key fields could be used as an alternate index
783    /// to speed up this filter.
784    UnusableNoSubset,
785}