mz_expr/relation/join_input_mapper.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
10use std::collections::BTreeSet;
11use std::ops::Range;
12
13use itertools::Itertools;
14use mz_repr::{ReprRelationType, SqlRelationType};
15
16use crate::visit::Visit;
17use crate::{MirRelationExpr, MirScalarExpr, VariadicFunc};
18
19/// Any column in a join expression exists in two contexts:
20/// 1) It has a position relative to the result of the join (global)
21/// 2) It has a position relative to the specific input it came from (local)
22/// This utility focuses on taking expressions that are in terms of
23/// the local input and re-expressing them in global terms and vice versa.
24///
25/// Methods in this class that take an argument `equivalences` are only
26/// guaranteed to return a correct answer if equivalence classes are in
27/// canonical form.
28/// (See [`crate::relation::canonicalize::canonicalize_equivalences`].)
29#[derive(Debug)]
30pub struct JoinInputMapper {
31 /// The number of columns per input. All other fields in this struct are
32 /// derived using the information in this field.
33 arities: Vec<usize>,
34 /// Looks up which input each column belongs to. Derived from `arities`.
35 /// Stored as a field to avoid recomputation.
36 input_relation: Vec<usize>,
37 /// The sum of the arities of the previous inputs in the join. Derived from
38 /// `arities`. Stored as a field to avoid recomputation.
39 prior_arities: Vec<usize>,
40}
41
42impl JoinInputMapper {
43 /// Creates a new `JoinInputMapper` and calculates the mapping of global context
44 /// columns to local context columns.
45 pub fn new(inputs: &[MirRelationExpr]) -> Self {
46 Self::new_from_input_arities(inputs.iter().map(|i| i.arity()))
47 }
48
49 /// Creates a new `JoinInputMapper` and calculates the mapping of global context
50 /// columns to local context columns. Using this method is more
51 /// efficient if input SQL types have been pre-calculated
52 pub fn new_from_input_types(types: &[SqlRelationType]) -> Self {
53 Self::new_from_input_arities(types.iter().map(|t| t.arity()))
54 }
55
56 /// Creates a new `JoinInputMapper` and calculates the mapping of global context
57 /// columns to local context columns. Using this method is more
58 /// efficient if input repr types have been pre-calculated.
59 pub fn new_from_input_repr_types(types: &[ReprRelationType]) -> Self {
60 Self::new_from_input_arities(types.iter().map(|t| t.arity()))
61 }
62
63 /// Creates a new `JoinInputMapper` and calculates the mapping of global context
64 /// columns to local context columns. Using this method is more
65 /// efficient if input arities have been pre-calculated
66 pub fn new_from_input_arities<I>(arities: I) -> Self
67 where
68 I: IntoIterator<Item = usize>,
69 {
70 let arities = arities.into_iter().collect::<Vec<usize>>();
71 let mut offset = 0;
72 let mut prior_arities = Vec::new();
73 for input in 0..arities.len() {
74 prior_arities.push(offset);
75 offset += arities[input];
76 }
77
78 let input_relation = arities
79 .iter()
80 .enumerate()
81 .flat_map(|(r, a)| std::iter::repeat(r).take(*a))
82 .collect::<Vec<_>>();
83
84 JoinInputMapper {
85 arities,
86 input_relation,
87 prior_arities,
88 }
89 }
90
91 /// reports sum of the number of columns of each input
92 pub fn total_columns(&self) -> usize {
93 self.arities.iter().sum()
94 }
95
96 /// reports total numbers of inputs in the join
97 pub fn total_inputs(&self) -> usize {
98 self.arities.len()
99 }
100
101 /// Using the keys that came from each local input,
102 /// figures out which keys remain unique in the larger join
103 /// Currently, we only figure out a small subset of the keys that
104 /// can remain unique.
105 pub fn global_keys<'a, I>(
106 &self,
107 mut local_keys: I,
108 equivalences: &[Vec<MirScalarExpr>],
109 ) -> Vec<Vec<usize>>
110 where
111 I: Iterator<Item = &'a Vec<Vec<usize>>>,
112 {
113 // A relation's uniqueness constraint holds if there is a
114 // sequence of the other relations such that each one has
115 // a uniqueness constraint whose columns are used in join
116 // constraints with relations prior in the sequence.
117 //
118 // Currently, we only:
119 // 1. test for whether the uniqueness constraints for the first input will hold
120 // 2. try one sequence, namely the inputs in order
121 // 3. check that the column themselves are used in the join constraints
122 // Technically uniqueness constraint would still hold if a 1-to-1
123 // expression on a unique key is used in the join constraint.
124
125 // for inputs `1..self.total_inputs()`, store a set of columns from that
126 // input that exist in join constraints that have expressions belonging to
127 // earlier inputs.
128 let mut column_with_prior_bound_by_input = vec![BTreeSet::new(); self.total_inputs() - 1];
129 for equivalence in equivalences {
130 // do a scan to find the first input represented in the constraint
131 let min_bound_input = equivalence
132 .iter()
133 .flat_map(|expr| self.lookup_inputs(expr).max())
134 .min();
135 if let Some(min_bound_input) = min_bound_input {
136 for expr in equivalence {
137 // then store all columns in the constraint that don't come
138 // from the first input
139 if let MirScalarExpr::Column(c, _name) = expr {
140 let (col, input) = self.map_column_to_local(*c);
141 if input > min_bound_input {
142 column_with_prior_bound_by_input[input - 1].insert(col);
143 }
144 }
145 }
146 }
147 }
148
149 if self.total_inputs() > 0 {
150 let first_input_keys = local_keys.next().unwrap().clone();
151 // for inputs `1..self.total_inputs()`, checks the keys belong to each
152 // input against the storage of columns that exist in join constraints
153 // that have expressions belonging to earlier inputs.
154 let remains_unique = local_keys.enumerate().all(|(index, keys)| {
155 keys.iter().any(|ks| {
156 ks.iter()
157 .all(|k| column_with_prior_bound_by_input[index].contains(k))
158 })
159 });
160
161 if remains_unique {
162 return first_input_keys;
163 }
164 }
165 vec![]
166 }
167
168 /// returns the arity for a particular input
169 #[inline]
170 pub fn input_arity(&self, index: usize) -> usize {
171 self.arities[index]
172 }
173
174 /// All column numbers in order for a particular input in the local context
175 #[inline]
176 pub fn local_columns(&self, index: usize) -> Range<usize> {
177 0..self.arities[index]
178 }
179
180 /// All column numbers in order for a particular input in the global context
181 #[inline]
182 pub fn global_columns(&self, index: usize) -> Range<usize> {
183 self.prior_arities[index]..(self.prior_arities[index] + self.arities[index])
184 }
185
186 /// Takes an expression from the global context and creates a new version
187 /// where column references have been remapped to the local context.
188 /// Assumes that all columns in `expr` are from the same input.
189 pub fn map_expr_to_local(&self, mut expr: MirScalarExpr) -> MirScalarExpr {
190 expr.visit_pre_mut(|e| {
191 if let MirScalarExpr::Column(c, _name) = e {
192 *c -= self.prior_arities[self.input_relation[*c]];
193 }
194 });
195 expr
196 }
197
198 /// Takes an expression from the local context of the `index`th input and
199 /// creates a new version where column references have been remapped to the
200 /// global context.
201 pub fn map_expr_to_global(&self, mut expr: MirScalarExpr, index: usize) -> MirScalarExpr {
202 expr.visit_pre_mut(|e| {
203 if let MirScalarExpr::Column(c, _name) = e {
204 *c += self.prior_arities[index];
205 }
206 });
207 expr
208 }
209
210 /// Remap column numbers from the global to the local context.
211 /// Returns a 2-tuple `(<new column number>, <index of input>)`
212 pub fn map_column_to_local(&self, column: usize) -> (usize, usize) {
213 let index = self.input_relation[column];
214 (column - self.prior_arities[index], index)
215 }
216
217 /// Remap a column number from a local context to the global context.
218 pub fn map_column_to_global(&self, column: usize, index: usize) -> usize {
219 column + self.prior_arities[index]
220 }
221
222 /// Takes a sequence of columns in the global context and splits it into
223 /// a `Vec` containing `self.total_inputs()` `BTreeSet`s, each containing
224 /// the localized columns that belong to the particular input.
225 pub fn split_column_set_by_input<'a, I>(&self, columns: I) -> Vec<BTreeSet<usize>>
226 where
227 I: Iterator<Item = &'a usize>,
228 {
229 let mut new_columns = vec![BTreeSet::new(); self.total_inputs()];
230 for column in columns {
231 let (new_col, input) = self.map_column_to_local(*column);
232 new_columns[input].extend(std::iter::once(new_col));
233 }
234 new_columns
235 }
236
237 /// Find the sorted, dedupped set of inputs an expression references
238 pub fn lookup_inputs(&self, expr: &MirScalarExpr) -> impl Iterator<Item = usize> + use<> {
239 expr.support()
240 .iter()
241 .map(|c| self.input_relation[*c])
242 .sorted()
243 .dedup()
244 }
245
246 /// Returns the index of the only input referenced in the given expression.
247 pub fn single_input(&self, expr: &MirScalarExpr) -> Option<usize> {
248 let mut inputs = self.lookup_inputs(expr);
249 if let Some(first_input) = inputs.next() {
250 if inputs.next().is_none() {
251 return Some(first_input);
252 }
253 }
254 None
255 }
256
257 /// Returns whether the given expr refers to columns of only the `index`th input.
258 pub fn is_localized(&self, expr: &MirScalarExpr, index: usize) -> bool {
259 if let Some(single_input) = self.single_input(expr) {
260 if single_input == index {
261 return true;
262 }
263 }
264 false
265 }
266
267 /// Takes an expression in the global context and looks in `equivalences`
268 /// for an equivalent expression (also expressed in the global context) that
269 /// belongs to one or more of the inputs in `bound_inputs`
270 ///
271 /// # Examples
272 ///
273 /// ```
274 /// use mz_repr::{Datum, SqlColumnType, SqlRelationType, SqlScalarType};
275 /// use mz_expr::{JoinInputMapper, MirRelationExpr, MirScalarExpr};
276 ///
277 /// // A two-column schema common to each of the three inputs
278 /// let schema = SqlRelationType::new(vec![
279 /// SqlScalarType::Int32.nullable(false),
280 /// SqlScalarType::Int32.nullable(false),
281 /// ]);
282 ///
283 /// // the specific data are not important here.
284 /// let data = vec![Datum::Int32(0), Datum::Int32(1)];
285 /// let input0 = MirRelationExpr::constant(vec![data.clone()], schema.clone());
286 /// let input1 = MirRelationExpr::constant(vec![data.clone()], schema.clone());
287 /// let input2 = MirRelationExpr::constant(vec![data.clone()], schema.clone());
288 ///
289 /// // [input0(#0) = input2(#1)], [input0(#1) = input1(#0) = input2(#0)]
290 /// let equivalences = vec![
291 /// vec![MirScalarExpr::column(0), MirScalarExpr::column(5)],
292 /// vec![MirScalarExpr::column(1), MirScalarExpr::column(2), MirScalarExpr::column(4)],
293 /// ];
294 ///
295 /// let input_mapper = JoinInputMapper::new(&[input0, input1, input2]);
296 /// assert_eq!(
297 /// Some(MirScalarExpr::column(4)),
298 /// input_mapper.find_bound_expr(&MirScalarExpr::column(2), &[2], &equivalences)
299 /// );
300 /// assert_eq!(
301 /// None,
302 /// input_mapper.find_bound_expr(&MirScalarExpr::column(0), &[1], &equivalences)
303 /// );
304 /// ```
305 pub fn find_bound_expr(
306 &self,
307 expr: &MirScalarExpr,
308 bound_inputs: &[usize],
309 equivalences: &[Vec<MirScalarExpr>],
310 ) -> Option<MirScalarExpr> {
311 if let Some(equivalence) = equivalences.iter().find(|equivs| equivs.contains(expr)) {
312 if let Some(bound_expr) = equivalence
313 .iter()
314 .find(|expr| self.lookup_inputs(expr).all(|i| bound_inputs.contains(&i)))
315 {
316 return Some(bound_expr.clone());
317 }
318 }
319 None
320 }
321
322 /// Try to rewrite `expr` from the global context so that all the
323 /// columns point to the `index`th input by replacing subexpressions with their
324 /// bound equivalents in the `index`th input if necessary.
325 /// Returns whether the rewriting was successful.
326 /// If it returns true, then `expr` is in the context of the `index`th input.
327 /// If it returns false, then still some subexpressions might have been rewritten. However,
328 /// `expr` is still in the global context.
329 pub fn try_localize_to_input_with_bound_expr(
330 &self,
331 expr: &mut MirScalarExpr,
332 index: usize,
333 equivalences: &[Vec<MirScalarExpr>],
334 ) -> bool {
335 // TODO (wangandi): Consider changing this code to be post-order
336 // instead of pre-order? `lookup_inputs` traverses all the nodes in
337 // `e` anyway, so we end up visiting nodes in `e` multiple times
338 // here. Alternatively, consider having the future `PredicateKnowledge`
339 // take over the responsibilities of this code?
340 #[allow(deprecated)]
341 expr.visit_mut_pre_post_nolimit(
342 &mut |e| {
343 let mut inputs = self.lookup_inputs(e);
344 if let Some(first_input) = inputs.next() {
345 if inputs.next().is_none() && first_input == index {
346 // there is only one input, and it is equal to index, so we're
347 // good. do not continue the recursion
348 return Some(vec![]);
349 }
350 }
351
352 if let Some(bound_expr) = self.find_bound_expr(e, &[index], equivalences) {
353 // Replace the subexpression with the equivalent one from input `index`
354 *e = bound_expr;
355 // The entire subexpression has been rewritten, so there is
356 // no need to visit any child expressions.
357 Some(vec![])
358 } else {
359 None
360 }
361 },
362 &mut |_| {},
363 );
364 if self.is_localized(expr, index) {
365 // If the localization attempt is successful, all columns in `expr`
366 // should only come from input `index`. Switch to the local context.
367 *expr = self.map_expr_to_local(expr.clone());
368 return true;
369 }
370 false
371 }
372
373 /// Try to find a consequence `c` of the given expression `e` for the given input.
374 ///
375 /// If we return `Some(c)`, that means
376 /// 1. `c` uses only columns from the given input;
377 /// 2. if `c` doesn't hold on a row of the input, then `e` also wouldn't hold;
378 /// 3. if `c` holds on a row of the input, then `e` might or might not hold.
379 /// 1. and 2. means that if we have a join with predicate `e` then we can use `c` for
380 /// pre-filtering a join input before the join. However, 3. means that `e` shouldn't be deleted
381 /// from the join predicates, i.e., we can't do a "traditional" predicate pushdown.
382 ///
383 /// Note that "`c` is a consequence of `e`" is the same thing as 2., see
384 /// <https://en.wikipedia.org/wiki/Contraposition>
385 ///
386 /// Example: For
387 /// `(t1.f2 = 3 AND t2.f2 = 4) OR (t1.f2 = 5 AND t2.f2 = 6)`
388 /// we find
389 /// `t1.f2 = 3 OR t1.f2 = 5` for t1, and
390 /// `t2.f2 = 4 OR t2.f2 = 6` for t2.
391 ///
392 /// Further examples are in TPC-H Q07, Q19, and chbench Q07, Q19.
393 ///
394 /// Parameters:
395 /// - `expr`: The expression `e` from above. `try_localize_to_input_with_bound_expr` should
396 /// be called on `expr` before us!
397 /// - `index`: The index of the join input whose columns we will use.
398 /// - `equivalences`: Join equivalences that we can use for `try_map_to_input_with_bound_expr`.
399 /// If successful, the returned expression is in the local context of the specified input.
400 pub fn consequence_for_input(
401 &self,
402 expr: &MirScalarExpr,
403 index: usize,
404 ) -> Option<MirScalarExpr> {
405 if self.is_localized(expr, index) {
406 Some(self.map_expr_to_local(expr.clone()))
407 } else {
408 match expr {
409 MirScalarExpr::CallVariadic {
410 func: VariadicFunc::Or,
411 exprs: or_args,
412 } => {
413 // Each OR arg should provide a consequence. If they do, we OR them.
414 let consequences_per_arg = or_args
415 .into_iter()
416 .map(|or_arg| {
417 mz_ore::stack::maybe_grow(|| self.consequence_for_input(or_arg, index))
418 })
419 .collect::<Option<Vec<_>>>()?; // return None if any of them are None
420 Some(MirScalarExpr::CallVariadic {
421 func: VariadicFunc::Or,
422 exprs: consequences_per_arg,
423 })
424 }
425 MirScalarExpr::CallVariadic {
426 func: VariadicFunc::And,
427 exprs: and_args,
428 } => {
429 // If any of the AND args provide a consequence, then we take those that do,
430 // and AND them.
431 let consequences_per_arg = and_args
432 .into_iter()
433 .map(|and_arg| {
434 mz_ore::stack::maybe_grow(|| self.consequence_for_input(and_arg, index))
435 })
436 .flat_map(|c| c) // take only those that provide a consequence
437 .collect_vec();
438 if consequences_per_arg.is_empty() {
439 None
440 } else {
441 Some(MirScalarExpr::CallVariadic {
442 func: VariadicFunc::And,
443 exprs: consequences_per_arg,
444 })
445 }
446 }
447 _ => None,
448 }
449 }
450 }
451}
452
453#[cfg(test)]
454mod tests {
455 use mz_repr::{Datum, SqlScalarType};
456
457 use crate::scalar::func;
458 use crate::{BinaryFunc, MirScalarExpr, UnaryFunc};
459
460 use super::*;
461
462 #[mz_ore::test]
463 #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `rust_psm_stack_pointer` on OS `linux`
464 fn try_map_to_input_with_bound_expr_test() {
465 let input_mapper = JoinInputMapper {
466 arities: vec![2, 3, 3],
467 input_relation: vec![0, 0, 1, 1, 1, 2, 2, 2],
468 prior_arities: vec![0, 2, 5],
469 };
470
471 // keys are numbered by (equivalence class #, input #)
472 let key10 = MirScalarExpr::column(0);
473 let key12 = MirScalarExpr::column(6);
474 let localized_key12 = MirScalarExpr::column(1);
475
476 let mut equivalences = vec![vec![key10.clone(), key12.clone()]];
477
478 // when the column is already part of the target input, all that happens
479 // is that it gets localized
480 let mut cloned = key12.clone();
481 input_mapper.try_localize_to_input_with_bound_expr(&mut cloned, 2, &equivalences);
482 assert_eq!(MirScalarExpr::column(1), cloned);
483
484 // basic tests that we can find a column's corresponding column in a
485 // different input
486 let mut cloned = key12.clone();
487 input_mapper.try_localize_to_input_with_bound_expr(&mut cloned, 0, &equivalences);
488 assert_eq!(key10, cloned);
489 let mut cloned = key12.clone();
490 assert_eq!(
491 false,
492 input_mapper.try_localize_to_input_with_bound_expr(&mut cloned, 1, &equivalences),
493 );
494
495 let key20 = MirScalarExpr::CallUnary {
496 func: UnaryFunc::NegInt32(crate::func::NegInt32),
497 expr: Box::new(MirScalarExpr::column(1)),
498 };
499 let key21 = MirScalarExpr::CallBinary {
500 func: BinaryFunc::AddInt32(func::AddInt32),
501 expr1: Box::new(MirScalarExpr::column(2)),
502 expr2: Box::new(MirScalarExpr::literal(
503 Ok(Datum::Int32(4)),
504 SqlScalarType::Int32,
505 )),
506 };
507 let key22 = MirScalarExpr::column(5);
508 let localized_key22 = MirScalarExpr::column(0);
509 equivalences.push(vec![key22.clone(), key20.clone(), key21.clone()]);
510
511 // basic tests that we can find an expression's corresponding expression in a
512 // different input
513 let mut cloned = key21.clone();
514 input_mapper.try_localize_to_input_with_bound_expr(&mut cloned, 0, &equivalences);
515 assert_eq!(key20, cloned);
516 let mut cloned = key21.clone();
517 input_mapper.try_localize_to_input_with_bound_expr(&mut cloned, 2, &equivalences);
518 assert_eq!(localized_key22, cloned);
519
520 // test that `try_map_to_input_with_bound_expr` will map multiple
521 // subexpressions to the corresponding expressions bound to a different input
522 let key_comp = MirScalarExpr::CallBinary {
523 func: func::MulInt32.into(),
524 expr1: Box::new(key12.clone()),
525 expr2: Box::new(key22),
526 };
527 let mut cloned = key_comp.clone();
528 input_mapper.try_localize_to_input_with_bound_expr(&mut cloned, 0, &equivalences);
529 assert_eq!(
530 MirScalarExpr::CallBinary {
531 func: func::MulInt32.into(),
532 expr1: Box::new(key10.clone()),
533 expr2: Box::new(key20.clone()),
534 },
535 cloned,
536 );
537
538 // test that the function returns None when part
539 // of the expression can be mapped to an input but the rest can't
540 let mut cloned = key_comp.clone();
541 assert_eq!(
542 false,
543 input_mapper.try_localize_to_input_with_bound_expr(&mut cloned, 1, &equivalences),
544 );
545
546 let key_comp_plus_non_key = MirScalarExpr::CallBinary {
547 func: func::Eq.into(),
548 expr1: Box::new(key_comp),
549 expr2: Box::new(MirScalarExpr::column(7)),
550 };
551 let mut mutab = key_comp_plus_non_key;
552 assert_eq!(
553 false,
554 input_mapper.try_localize_to_input_with_bound_expr(&mut mutab, 0, &equivalences),
555 );
556
557 let key_comp_multi_input = MirScalarExpr::CallBinary {
558 func: func::Eq.into(),
559 expr1: Box::new(key12),
560 expr2: Box::new(key21),
561 };
562 // test that the function works when part of the expression is already
563 // part of the target input
564 let mut cloned = key_comp_multi_input.clone();
565 input_mapper.try_localize_to_input_with_bound_expr(&mut cloned, 2, &equivalences);
566 assert_eq!(
567 MirScalarExpr::CallBinary {
568 func: func::Eq.into(),
569 expr1: Box::new(localized_key12),
570 expr2: Box::new(localized_key22),
571 },
572 cloned,
573 );
574 // test that the function works when parts of the expression come from
575 // multiple inputs
576 let mut cloned = key_comp_multi_input.clone();
577 input_mapper.try_localize_to_input_with_bound_expr(&mut cloned, 0, &equivalences);
578 assert_eq!(
579 MirScalarExpr::CallBinary {
580 func: func::Eq.into(),
581 expr1: Box::new(key10),
582 expr2: Box::new(key20),
583 },
584 cloned,
585 );
586 let mut mutab = key_comp_multi_input;
587 assert_eq!(
588 false,
589 input_mapper.try_localize_to_input_with_bound_expr(&mut mutab, 1, &equivalences),
590 )
591 }
592}