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