mz_expr/linear.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.
9use std::collections::{BTreeMap, BTreeSet};
10use std::fmt::Display;
11
12use mz_repr::{Datum, Row};
13use serde::{Deserialize, Serialize};
14
15use crate::scalar::optimizable::OptimizableExpr;
16use crate::visit::Visit;
17use crate::{MirRelationExpr, MirScalarExpr};
18
19/// A compound operator that can be applied row-by-row.
20///
21/// This operator integrates the map, filter, and project operators.
22/// It applies a sequences of map expressions, which are allowed to
23/// refer to previous expressions, interleaved with predicates which
24/// must be satisfied for an output to be produced. If all predicates
25/// evaluate to `Datum::True` the data at the identified columns are
26/// collected and produced as output in a packed `Row`.
27///
28/// This operator is a "builder" and its contents may contain expressions
29/// that are not yet executable. For example, it may contain temporal
30/// expressions in `self.expressions`, even though this is not something
31/// we can directly evaluate. The plan creation methods will defensively
32/// ensure that the right thing happens.
33#[derive(
34 Clone,
35 Debug,
36 Eq,
37 PartialEq,
38 Serialize,
39 Deserialize,
40 Hash,
41 Ord,
42 PartialOrd
43)]
44#[serde(bound(deserialize = "E: serde::de::DeserializeOwned"))]
45pub struct MapFilterProject<E: OptimizableExpr = MirScalarExpr> {
46 /// A sequence of expressions that should be appended to the row.
47 ///
48 /// Many of these expressions may not be produced in the output,
49 /// and may only be present as common subexpressions.
50 pub expressions: Vec<E>,
51 /// Expressions that must evaluate to `Datum::True` for the output
52 /// row to be produced.
53 ///
54 /// Each entry is prepended with a column identifier indicating
55 /// the column *before* which the predicate should first be applied.
56 /// Most commonly this would be one plus the largest column identifier
57 /// in the predicate's support, but it could be larger to implement
58 /// guarded evaluation of predicates.
59 ///
60 /// This list should be sorted by the first field.
61 pub predicates: Vec<(usize, E)>,
62 /// A sequence of column identifiers whose data form the output row.
63 pub projection: Vec<usize>,
64 /// The expected number of input columns.
65 ///
66 /// This is needed to ensure correct identification of newly formed
67 /// columns in the output.
68 pub input_arity: usize,
69}
70
71impl<E: OptimizableExpr + Display> Display for MapFilterProject<E> {
72 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
73 writeln!(f, "MapFilterProject(")?;
74 writeln!(f, " expressions:")?;
75 self.expressions
76 .iter()
77 .enumerate()
78 .try_for_each(|(i, e)| writeln!(f, " #{} <- {},", i + self.input_arity, e))?;
79 writeln!(f, " predicates:")?;
80 self.predicates
81 .iter()
82 .try_for_each(|(before, p)| writeln!(f, " <before: {}> {},", before, p))?;
83 writeln!(f, " projection: {:?}", self.projection)?;
84 writeln!(f, " input_arity: {}", self.input_arity)?;
85 writeln!(f, ")")
86 }
87}
88
89impl<E: OptimizableExpr> MapFilterProject<E> {
90 /// Create a no-op operator for an input of a supplied arity.
91 pub fn new(input_arity: usize) -> Self {
92 Self {
93 expressions: Vec::new(),
94 predicates: Vec::new(),
95 projection: (0..input_arity).collect(),
96 input_arity,
97 }
98 }
99
100 /// Given two mfps, return an mfp that applies one
101 /// followed by the other.
102 /// Note that the arguments are in the opposite order
103 /// from how function composition is usually written in mathematics.
104 pub fn compose(before: Self, after: Self) -> Self {
105 let (m, f, p) = after.into_map_filter_project();
106 before.map(m).filter(f).project(p)
107 }
108
109 /// True if the operator describes the identity transformation.
110 pub fn is_identity(&self) -> bool {
111 self.expressions.is_empty()
112 && self.predicates.is_empty()
113 && self.projection.len() == self.input_arity
114 && self.projection.iter().enumerate().all(|(i, p)| i == *p)
115 }
116
117 /// Retain only the indicated columns in the presented order.
118 pub fn project<I>(mut self, columns: I) -> Self
119 where
120 I: IntoIterator<Item = usize> + std::fmt::Debug,
121 {
122 self.projection = columns.into_iter().map(|c| self.projection[c]).collect();
123 self
124 }
125
126 /// Retain only rows satisfying these predicates.
127 ///
128 /// This method introduces predicates as eagerly as they can be evaluated,
129 /// which may not be desired for predicates that may cause exceptions.
130 /// If fine manipulation is required, the predicates can be added manually.
131 pub fn filter<I>(mut self, predicates: I) -> Self
132 where
133 I: IntoIterator<Item = E>,
134 {
135 for mut predicate in predicates {
136 // Correct column references.
137 predicate.permute(&self.projection[..]);
138
139 // Validate column references.
140 assert!(
141 predicate
142 .support()
143 .into_iter()
144 .all(|c| c < self.input_arity + self.expressions.len())
145 );
146
147 // Insert predicate as eagerly as it can be evaluated:
148 // just after the largest column in its support is formed.
149 let max_support = predicate
150 .support()
151 .into_iter()
152 .max()
153 .map(|c| c + 1)
154 .unwrap_or(0);
155 self.predicates.push((max_support, predicate))
156 }
157 // Stable sort predicates by position at which they take effect.
158 // We put literal errors at the end as a stop-gap to avoid erroring
159 // before we are able to evaluate any predicates that might prevent it.
160 self.predicates
161 .sort_by_key(|(position, predicate)| (E::is_literal_err(predicate), *position));
162 self
163 }
164
165 /// Append the result of evaluating expressions to each row.
166 pub fn map<I>(mut self, expressions: I) -> Self
167 where
168 I: IntoIterator<Item = E>,
169 {
170 for mut expression in expressions {
171 // Correct column references.
172 expression.permute(&self.projection[..]);
173
174 // Validate column references.
175 assert!(
176 expression
177 .support()
178 .into_iter()
179 .all(|c| c < self.input_arity + self.expressions.len())
180 );
181
182 // Introduce expression and produce as output.
183 self.expressions.push(expression);
184 self.projection
185 .push(self.input_arity + self.expressions.len() - 1);
186 }
187
188 self
189 }
190
191 /// Like [`MapFilterProject::as_map_filter_project`], but consumes `self` rather than cloning.
192 pub fn into_map_filter_project(self) -> (Vec<E>, Vec<E>, Vec<usize>) {
193 let predicates = self
194 .predicates
195 .into_iter()
196 .map(|(_pos, predicate)| predicate)
197 .collect();
198 (self.expressions, predicates, self.projection)
199 }
200
201 /// As the arguments to `Map`, `Filter`, and `Project` operators.
202 ///
203 /// In principle, this operator can be implemented as a sequence of
204 /// more elemental operators, likely less efficiently.
205 pub fn as_map_filter_project(&self) -> (Vec<E>, Vec<E>, Vec<usize>) {
206 self.clone().into_map_filter_project()
207 }
208}
209
210// Methods that are specific to MirScalarExpr (use MirRelationExpr or as_literal).
211impl MapFilterProject<MirScalarExpr> {
212 /// Determines if a scalar expression must be equal to a literal datum.
213 pub fn literal_constraint(&self, expr: &MirScalarExpr) -> Option<Datum<'_>> {
214 for (_pos, predicate) in self.predicates.iter() {
215 if let MirScalarExpr::CallBinary {
216 func: crate::BinaryFunc::Eq(_),
217 expr1,
218 expr2,
219 } = predicate
220 {
221 if let Some(Ok(datum1)) = expr1.as_literal() {
222 if &**expr2 == expr {
223 return Some(datum1);
224 }
225 }
226 if let Some(Ok(datum2)) = expr2.as_literal() {
227 if &**expr1 == expr {
228 return Some(datum2);
229 }
230 }
231 }
232 }
233 None
234 }
235
236 /// Determines if a sequence of scalar expressions must be equal to a literal row.
237 ///
238 /// This method returns `None` on an empty `exprs`, which might be surprising, but
239 /// seems to line up with its callers' expectations of that being a non-constraint.
240 /// The caller knows if `exprs` is empty, and can modify their behavior appropriately.
241 /// if they would rather have a literal empty row.
242 pub fn literal_constraints(&self, exprs: &[MirScalarExpr]) -> Option<Row> {
243 if exprs.is_empty() {
244 return None;
245 }
246 let mut row = Row::default();
247 let mut packer = row.packer();
248 for expr in exprs {
249 if let Some(literal) = self.literal_constraint(expr) {
250 packer.push(literal);
251 } else {
252 return None;
253 }
254 }
255 Some(row)
256 }
257
258 /// Extracts any MapFilterProject at the root of the expression.
259 ///
260 /// The expression will be modified to extract any maps, filters, and
261 /// projections, which will be returned as `Self`. If there are no maps,
262 /// filters, or projections the method will return an identity operator.
263 ///
264 /// The extracted expressions may contain temporal predicates, and one
265 /// should be careful to apply them blindly.
266 pub fn extract_from_expression(expr: &MirRelationExpr) -> (Self, &MirRelationExpr) {
267 // TODO: This could become iterative rather than recursive if
268 // we were able to fuse MFP operators from below, rather than
269 // from above.
270 match expr {
271 MirRelationExpr::Map { input, scalars } => {
272 let (mfp, expr) = Self::extract_from_expression(input);
273 (mfp.map(scalars.iter().cloned()), expr)
274 }
275 MirRelationExpr::Filter { input, predicates } => {
276 let (mfp, expr) = Self::extract_from_expression(input);
277 (mfp.filter(predicates.iter().cloned()), expr)
278 }
279 MirRelationExpr::Project { input, outputs } => {
280 let (mfp, expr) = Self::extract_from_expression(input);
281 (mfp.project(outputs.iter().cloned()), expr)
282 }
283 // TODO: The recursion is quadratic in the number of Map/Filter/Project operators due to
284 // this call to `arity()`.
285 x => (Self::new(x.arity()), x),
286 }
287 }
288
289 /// Extracts an error-free MapFilterProject at the root of the expression.
290 ///
291 /// The expression will be modified to extract maps, filters, and projects
292 /// from the root of the expression, which will be returned as `Self`. The
293 /// extraction will halt if a Map or Filter containing a literal error is
294 /// reached. Otherwise, the method will return an identity operator.
295 ///
296 /// This method is meant to be used during optimization, where it is
297 /// necessary to avoid moving around maps and filters with errors.
298 pub fn extract_non_errors_from_expr(expr: &MirRelationExpr) -> (Self, &MirRelationExpr) {
299 match expr {
300 MirRelationExpr::Map { input, scalars }
301 if scalars.iter().all(|s| !s.is_literal_err()) =>
302 {
303 let (mfp, expr) = Self::extract_non_errors_from_expr(input);
304 (mfp.map(scalars.iter().cloned()), expr)
305 }
306 MirRelationExpr::Filter { input, predicates }
307 if predicates.iter().all(|p| !p.is_literal_err()) =>
308 {
309 let (mfp, expr) = Self::extract_non_errors_from_expr(input);
310 (mfp.filter(predicates.iter().cloned()), expr)
311 }
312 MirRelationExpr::Project { input, outputs } => {
313 let (mfp, expr) = Self::extract_non_errors_from_expr(input);
314 (mfp.project(outputs.iter().cloned()), expr)
315 }
316 x => (Self::new(x.arity()), x),
317 }
318 }
319
320 /// Extracts an error-free MapFilterProject at the root of the expression.
321 ///
322 /// Differs from [MapFilterProject::extract_non_errors_from_expr] by taking and returning a
323 /// mutable reference.
324 pub fn extract_non_errors_from_expr_ref_mut(
325 expr: &mut MirRelationExpr,
326 ) -> (Self, &mut MirRelationExpr) {
327 // This is essentially the same code as `extract_non_errors_from_expr`, except the seemingly
328 // superfluous outer if, which works around a borrow-checker issue:
329 // https://github.com/rust-lang/rust/issues/54663
330 if matches!(
331 expr,
332 MirRelationExpr::Map { input: _, scalars }
333 if scalars.iter().all(|s| !s.is_literal_err())
334 ) || matches!(
335 expr,
336 MirRelationExpr::Filter { input: _, predicates }
337 if predicates.iter().all(|p| !p.is_literal_err())
338 ) || matches!(expr, MirRelationExpr::Project { .. })
339 {
340 match expr {
341 MirRelationExpr::Map { input, scalars }
342 if scalars.iter().all(|s| !s.is_literal_err()) =>
343 {
344 let (mfp, expr) = Self::extract_non_errors_from_expr_ref_mut(input);
345 (mfp.map(scalars.iter().cloned()), expr)
346 }
347 MirRelationExpr::Filter { input, predicates }
348 if predicates.iter().all(|p| !p.is_literal_err()) =>
349 {
350 let (mfp, expr) = Self::extract_non_errors_from_expr_ref_mut(input);
351 (mfp.filter(predicates.iter().cloned()), expr)
352 }
353 MirRelationExpr::Project { input, outputs } => {
354 let (mfp, expr) = Self::extract_non_errors_from_expr_ref_mut(input);
355 (mfp.project(outputs.iter().cloned()), expr)
356 }
357 _ => unreachable!(),
358 }
359 } else {
360 (Self::new(expr.arity()), expr)
361 }
362 }
363
364 /// Removes an error-free MapFilterProject from the root of the expression.
365 ///
366 /// The expression will be modified to extract maps, filters, and projects
367 /// from the root of the expression, which will be returned as `Self`. The
368 /// extraction will halt if a Map or Filter containing a literal error is
369 /// reached. Otherwise, the method will return an
370 /// identity operator, and the expression will remain unchanged.
371 ///
372 /// This method is meant to be used during optimization, where it is
373 /// necessary to avoid moving around maps and filters with errors.
374 pub fn extract_non_errors_from_expr_mut(expr: &mut MirRelationExpr) -> Self {
375 match expr {
376 MirRelationExpr::Map { input, scalars }
377 if scalars.iter().all(|s| !s.is_literal_err()) =>
378 {
379 let mfp =
380 Self::extract_non_errors_from_expr_mut(input).map(scalars.iter().cloned());
381 *expr = input.take_dangerous();
382 mfp
383 }
384 MirRelationExpr::Filter { input, predicates }
385 if predicates.iter().all(|p| !p.is_literal_err()) =>
386 {
387 let mfp = Self::extract_non_errors_from_expr_mut(input)
388 .filter(predicates.iter().cloned());
389 *expr = input.take_dangerous();
390 mfp
391 }
392 MirRelationExpr::Project { input, outputs } => {
393 let mfp =
394 Self::extract_non_errors_from_expr_mut(input).project(outputs.iter().cloned());
395 *expr = input.take_dangerous();
396 mfp
397 }
398 x => Self::new(x.arity()),
399 }
400 }
401}
402
403impl<E: OptimizableExpr> MapFilterProject<E> {
404 /// Returns `true` if any predicate in this MFP contains a temporal expression (`mz_now()`).
405 pub fn has_temporal_predicates(&self) -> bool {
406 self.predicates
407 .iter()
408 .any(|(_, predicate)| OptimizableExpr::contains_temporal(predicate))
409 }
410
411 /// Extracts temporal predicates into their own `Self`.
412 ///
413 /// Expressions that are used by the temporal predicates are exposed by `self.projection`,
414 /// though there could be justification for extracting them as well if they are otherwise
415 /// unused.
416 ///
417 /// This separation is valuable when the execution cannot be fused into one operator.
418 pub fn extract_temporal(&mut self) -> Self {
419 // Optimize the expression, as it is only post-optimization that we can be certain
420 // that temporal expressions are restricted to filters. We could relax this in the
421 // future to be only `inline_expressions` and `remove_undemanded`, but optimization
422 // seems to be the best fit at the moment.
423 self.optimize();
424
425 // Assert that we no longer have temporal expressions to evaluate. This should only
426 // occur if the optimization above results with temporal expressions yielded in the
427 // output, which is out of spec for how the type is meant to be used.
428 assert!(
429 !self
430 .expressions
431 .iter()
432 .any(|e| OptimizableExpr::contains_temporal(e))
433 );
434
435 // Extract temporal predicates from `self.predicates`.
436 let mut temporal_predicates = Vec::new();
437 self.predicates.retain(|(_position, predicate)| {
438 if OptimizableExpr::contains_temporal(predicate) {
439 temporal_predicates.push(predicate.clone());
440 false
441 } else {
442 true
443 }
444 });
445
446 // Determine extended input columns used by temporal filters.
447 let mut support = BTreeSet::new();
448 for predicate in temporal_predicates.iter() {
449 support.extend(predicate.support());
450 }
451
452 // Discover the locations of these columns after `self.projection`.
453 let old_projection_len = self.projection.len();
454 let mut new_location = BTreeMap::new();
455 for original in support.iter() {
456 if let Some(position) = self.projection.iter().position(|x| x == original) {
457 new_location.insert(*original, position);
458 } else {
459 new_location.insert(*original, self.projection.len());
460 self.projection.push(*original);
461 }
462 }
463 // Permute references in extracted predicates to their new locations.
464 for predicate in temporal_predicates.iter_mut() {
465 predicate.permute_map(&new_location);
466 }
467
468 // Form a new `Self` containing the temporal predicates to return.
469 Self::new(self.projection.len())
470 .filter(temporal_predicates)
471 .project(0..old_projection_len)
472 }
473
474 /// Extracts common expressions from multiple `Self` into a result `Self`.
475 ///
476 /// The argument `mfps` are mutated so that each are functionaly equivalent to their
477 /// corresponding input, when composed atop the resulting `Self`.
478 ///
479 /// The `extract_exprs` argument is temporary, as we roll out the `extract_common_mfp_expressions` flag.
480 pub fn extract_common(mfps: &mut [&mut Self]) -> Self {
481 match mfps.len() {
482 0 => {
483 panic!("Cannot call method on empty arguments");
484 }
485 1 => {
486 let output_arity = mfps[0].projection.len();
487 std::mem::replace(mfps[0], MapFilterProject::new(output_arity))
488 }
489 _ => {
490 // More generally, we convert each mfp to ANF, at which point we can
491 // repeatedly extract atomic expressions that depend only on input
492 // columns, migrate them to an input mfp, and repeat until no such
493 // expressions exist. At this point, we can also migrate predicates
494 // and then determine and push down projections.
495
496 // Prepare a return `Self`.
497 let mut result_mfp = MapFilterProject::new(mfps[0].input_arity);
498
499 // We convert each mfp to ANF, using `memoize_expressions`.
500 for mfp in mfps.iter_mut() {
501 mfp.memoize_expressions();
502 }
503
504 // We repeatedly extract common expressions, until none remain.
505 let mut done = false;
506 while !done {
507 // We use references to determine common expressions, and must
508 // introduce a scope here to drop the borrows before mutation.
509 let common = {
510 // The input arity may increase as we iterate, so recapture.
511 let input_arity = result_mfp.projection.len();
512 let mut prev: BTreeSet<_> = mfps[0]
513 .expressions
514 .iter()
515 .filter(|e| e.support().last() < Some(&input_arity))
516 .collect();
517 let mut next = BTreeSet::default();
518 for mfp in mfps[1..].iter() {
519 for expr in mfp.expressions.iter() {
520 if prev.contains(expr) {
521 next.insert(expr);
522 }
523 }
524 std::mem::swap(&mut prev, &mut next);
525 next.clear();
526 }
527 prev.into_iter().cloned().collect::<Vec<_>>()
528 };
529 // Without new common expressions, we should terminate the loop.
530 done = common.is_empty();
531
532 // Migrate each expression in `common` to `result_mfp`.
533 for expr in common.into_iter() {
534 // Update each mfp by removing expr and updating column references.
535 for mfp in mfps.iter_mut() {
536 // With `expr` next in `result_mfp`, it is as if we are rotating it to
537 // be the first expression in `mfp`, and then removing it from `mfp` and
538 // increasing the input arity of `mfp`.
539 let arity = result_mfp.projection.len();
540 let found = mfp.expressions.iter().position(|e| e == &expr).unwrap();
541 let index = arity + found;
542 // Column references change due to the rotation from `index` to `arity`.
543 let action = |c: &mut usize| {
544 if arity <= *c && *c < index {
545 *c += 1;
546 } else if *c == index {
547 *c = arity;
548 }
549 };
550 // Rotate `expr` from `found` to first, and then snip.
551 // Short circuit by simply removing and incrementing the input arity.
552 mfp.input_arity += 1;
553 mfp.expressions.remove(found);
554 // Update column references in expressions, predicates, and projections.
555 for e in mfp.expressions.iter_mut() {
556 e.visit_columns(action);
557 }
558 for (o, e) in mfp.predicates.iter_mut() {
559 e.visit_columns(action);
560 // Max out the offset for the predicate; optimization will correct.
561 *o = mfp.input_arity + mfp.expressions.len();
562 }
563 for c in mfp.projection.iter_mut() {
564 action(c);
565 }
566 }
567 // Install the expression and update
568 result_mfp.expressions.push(expr);
569 result_mfp.projection.push(result_mfp.projection.len());
570 }
571 }
572 // As before, but easier: predicates in common to all mfps.
573 let common_preds: Vec<E> = {
574 let input_arity = result_mfp.projection.len();
575 let mut prev: BTreeSet<_> = mfps[0]
576 .predicates
577 .iter()
578 .map(|(_, e)| e)
579 .filter(|e| e.support().last() < Some(&input_arity))
580 .collect();
581 let mut next = BTreeSet::default();
582 for mfp in mfps[1..].iter() {
583 for (_, expr) in mfp.predicates.iter() {
584 if prev.contains(expr) {
585 next.insert(expr);
586 }
587 }
588 std::mem::swap(&mut prev, &mut next);
589 next.clear();
590 }
591 // Expressions in common, that we will append to `result_mfp.expressions`.
592 prev.into_iter().cloned().collect::<Vec<_>>()
593 };
594 for mfp in mfps.iter_mut() {
595 mfp.predicates.retain(|(_, p)| !common_preds.contains(p));
596 mfp.optimize();
597 }
598 result_mfp.predicates.extend(
599 common_preds
600 .into_iter()
601 .map(|e| (result_mfp.projection.len(), e)),
602 );
603
604 // Then, look for unused columns and project them away.
605 let mut common_demand = BTreeSet::new();
606 for mfp in mfps.iter() {
607 common_demand.extend(mfp.demand());
608 }
609 // columns in `common_demand` must be retained, but others
610 // may be discarded.
611 let common_demand = (0..result_mfp.projection.len())
612 .filter(|x| common_demand.contains(x))
613 .collect::<Vec<_>>();
614 let remap = common_demand
615 .iter()
616 .cloned()
617 .enumerate()
618 .map(|(new, old)| (old, new))
619 .collect::<BTreeMap<_, _>>();
620 for mfp in mfps.iter_mut() {
621 mfp.permute_fn(|c| remap[&c], common_demand.len());
622 }
623 result_mfp = result_mfp.project(common_demand);
624
625 // Return the resulting MFP.
626 result_mfp.optimize();
627 result_mfp
628 }
629 }
630 }
631
632 /// Returns `self`, and leaves behind an identity operator that acts on its output.
633 pub fn take(&mut self) -> Self {
634 let mut identity = Self::new(self.projection.len());
635 std::mem::swap(self, &mut identity);
636 identity
637 }
638
639 /// Convert the `MapFilterProject` into a staged evaluation plan.
640 ///
641 /// The main behavior is extract temporal predicates, which cannot be evaluated
642 /// using the standard machinery.
643 pub fn into_plan(self) -> Result<plan::MfpPlan<E>, String> {
644 plan::MfpPlan::create_from(self)
645 }
646}
647
648impl<E: OptimizableExpr> MapFilterProject<E> {
649 /// Partitions `self` into two instances, one of which can be eagerly applied.
650 ///
651 /// The `available` argument indicates which input columns are available (keys)
652 /// and in which positions (values). This information may allow some maps and
653 /// filters to execute. The `input_arity` argument reports the total number of
654 /// input columns (which may include some not present in `available`)
655 ///
656 /// This method partitions `self` in two parts, `(before, after)`, where `before`
657 /// can be applied on columns present as keys in `available`, and `after` must
658 /// await the introduction of the other input columns.
659 ///
660 /// The `before` instance will *append* any columns that can be determined from
661 /// `available` but will project away any of these columns that are not needed by
662 /// `after`. Importantly, this means that `before` will leave intact *all* input
663 /// columns including those not referenced in `available`.
664 ///
665 /// The `after` instance will presume all input columns are available, followed
666 /// by the appended columns of the `before` instance. It may be that some input
667 /// columns can be projected away in `before` if `after` does not need them, but
668 /// we leave that as something the caller can apply if needed (it is otherwise
669 /// complicated to negotiate which input columns `before` should retain).
670 ///
671 /// To correctly reconstruct `self` from `before` and `after`, one must introduce
672 /// additional input columns, permute all input columns to their locations as
673 /// expected by `self`, follow this by new columns appended by `before`, and
674 /// remove all other columns that may be present.
675 ///
676 /// # Example
677 ///
678 /// ```rust
679 /// use mz_expr::{BinaryFunc, MapFilterProject, MirScalarExpr, func};
680 ///
681 /// // imagine an action on columns (a, b, c, d).
682 /// let original = MapFilterProject::new(4).map(vec![
683 /// MirScalarExpr::column(0).call_binary(MirScalarExpr::column(1), func::AddInt64),
684 /// MirScalarExpr::column(2).call_binary(MirScalarExpr::column(4), func::AddInt64),
685 /// MirScalarExpr::column(3).call_binary(MirScalarExpr::column(5), func::AddInt64),
686 /// ]).project(vec![6]);
687 ///
688 /// // Imagine we start with columns (b, x, a, y, c).
689 /// //
690 /// // The `partition` method requires a map from *expected* input columns to *actual*
691 /// // input columns. In the example above, the columns a, b, and c exist, and are at
692 /// // locations 2, 0, and 4 respectively. We must construct a map to this effect.
693 /// let mut available_columns = std::collections::BTreeMap::new();
694 /// available_columns.insert(0, 2);
695 /// available_columns.insert(1, 0);
696 /// available_columns.insert(2, 4);
697 /// // Partition `original` using the available columns and current input arity.
698 /// // This informs `partition` which columns are available, where they can be found,
699 /// // and how many columns are not relevant but should be preserved.
700 /// let (before, after) = original.partition(available_columns, 5);
701 ///
702 /// // `before` sees all five input columns, and should append `a + b + c`.
703 /// assert_eq!(before, MapFilterProject::new(5).map(vec![
704 /// MirScalarExpr::column(2).call_binary(MirScalarExpr::column(0), func::AddInt64),
705 /// MirScalarExpr::column(4).call_binary(MirScalarExpr::column(5), func::AddInt64),
706 /// ]).project(vec![0, 1, 2, 3, 4, 6]));
707 ///
708 /// // `after` expects to see `(a, b, c, d, a + b + c)`.
709 /// assert_eq!(after, MapFilterProject::new(5).map(vec![
710 /// MirScalarExpr::column(3).call_binary(MirScalarExpr::column(4), func::AddInt64)
711 /// ]).project(vec![5]));
712 ///
713 /// // To reconstruct `self`, we must introduce the columns that are not present,
714 /// // and present them in the order intended by `self`. In this example, we must
715 /// // introduce column d and permute the columns so that they begin (a, b, c, d).
716 /// // The columns x and y must be projected away, and any columns introduced by
717 /// // `begin` must be retained in their current order.
718 ///
719 /// // The `after` instance expects to be provided with all inputs, but it
720 /// // may not need all inputs. The `demand()` and `permute()` methods can
721 /// // optimize the representation.
722 /// ```
723 pub fn partition(self, available: BTreeMap<usize, usize>, input_arity: usize) -> (Self, Self) {
724 // Map expressions, filter predicates, and projections for `before` and `after`.
725 let mut before_expr = Vec::new();
726 let mut before_pred = Vec::new();
727 let mut before_proj = Vec::new();
728 let mut after_expr = Vec::new();
729 let mut after_pred = Vec::new();
730 let mut after_proj = Vec::new();
731
732 // Track which output columns must be preserved in the output of `before`.
733 let mut demanded = BTreeSet::new();
734 demanded.extend(0..self.input_arity);
735 demanded.extend(self.projection.iter());
736
737 // Determine which map expressions can be computed from the available subset.
738 // Some expressions may depend on other expressions, but by evaluating them
739 // in forward order we should accurately determine the available expressions.
740 let mut available_expr = vec![false; self.input_arity];
741 // Initialize available columns from `available`, which is then not used again.
742 for index in available.keys() {
743 available_expr[*index] = true;
744 }
745 for expr in self.expressions.into_iter() {
746 // We treat an expression as available if its supporting columns are available,
747 // and if it is not a literal (we want to avoid pushing down literals). This
748 // choice is ad-hoc, but the intent is that we partition the operators so
749 // that we can reduce the row representation size and total computation.
750 // Pushing down literals harms the former and does nothing for the latter.
751 // In the future, we'll want to have a harder think about this trade-off, as
752 // we are certainly making sub-optimal decisions by pushing down all available
753 // work.
754 // TODO(mcsherry): establish better principles about what work to push down.
755 let is_available = expr.support().into_iter().all(|i| available_expr[i])
756 && !OptimizableExpr::is_literal(&expr);
757 if is_available {
758 before_expr.push(expr);
759 } else {
760 demanded.extend(expr.support());
761 after_expr.push(expr);
762 }
763 available_expr.push(is_available);
764 }
765
766 // Determine which predicates can be computed from the available subset.
767 for (_when, pred) in self.predicates.into_iter() {
768 let is_available = pred.support().into_iter().all(|i| available_expr[i]);
769 if is_available {
770 before_pred.push(pred);
771 } else {
772 demanded.extend(pred.support());
773 after_pred.push(pred);
774 }
775 }
776
777 // Map from prior output location to location in un-projected `before`.
778 // This map is used to correct references in `before` but it should be
779 // adjusted to reflect `before`s projection prior to use in `after`.
780 let mut before_map = available;
781 // Input columns include any additional undescribed columns that may
782 // not be captured by the `available` argument, so we must independently
783 // track the current number of columns (vs relying on `before_map.len()`).
784 let mut input_columns = input_arity;
785 for index in self.input_arity..available_expr.len() {
786 if available_expr[index] {
787 before_map.insert(index, input_columns);
788 input_columns += 1;
789 }
790 }
791
792 // Permute the column references in `before` expressions and predicates.
793 for expr in before_expr.iter_mut() {
794 expr.permute_map(&before_map);
795 }
796 for pred in before_pred.iter_mut() {
797 pred.permute_map(&before_map);
798 }
799
800 // Demand information determines `before`s output projection.
801 // Specifically, we produce all input columns in the output, as well as
802 // any columns that are available and demanded.
803 before_proj.extend(0..input_arity);
804 for index in self.input_arity..available_expr.len() {
805 // If an intermediate result is both available and demanded,
806 // we should produce it as output.
807 if available_expr[index] && demanded.contains(&index) {
808 // Use the new location of `index`.
809 before_proj.push(before_map[&index]);
810 }
811 }
812
813 // Map from prior output locations to location in post-`before` columns.
814 // This map is used to correct references in `after`.
815 // The presumption is that `after` will be presented with all input columns,
816 // followed by the output columns introduced by `before` in order.
817 let mut after_map = BTreeMap::new();
818 for index in 0..self.input_arity {
819 after_map.insert(index, index);
820 }
821 for index in self.input_arity..available_expr.len() {
822 // If an intermediate result is both available and demanded,
823 // it was produced as output.
824 if available_expr[index] && demanded.contains(&index) {
825 // We expect to find the output as far after `self.input_arity` as
826 // it was produced after `input_arity` in the output of `before`.
827 let location = self.input_arity
828 + (before_proj
829 .iter()
830 .position(|x| x == &before_map[&index])
831 .unwrap()
832 - input_arity);
833 after_map.insert(index, location);
834 }
835 }
836 // We must now re-map the remaining non-demanded expressions, which are
837 // contiguous rather than potentially interspersed.
838 for index in self.input_arity..available_expr.len() {
839 if !available_expr[index] {
840 after_map.insert(index, after_map.len());
841 }
842 }
843
844 // Permute the column references in `after` expressions and predicates.
845 for expr in after_expr.iter_mut() {
846 expr.permute_map(&after_map);
847 }
848 for pred in after_pred.iter_mut() {
849 pred.permute_map(&after_map);
850 }
851 // Populate `after` projection with the new locations of `self.projection`.
852 for index in self.projection {
853 after_proj.push(after_map[&index]);
854 }
855
856 // Form and return the before and after MapFilterProject instances.
857 let before = Self::new(input_arity)
858 .map(before_expr)
859 .filter(before_pred)
860 .project(before_proj.clone());
861 let after = Self::new(self.input_arity + (before_proj.len() - input_arity))
862 .map(after_expr)
863 .filter(after_pred)
864 .project(after_proj);
865 (before, after)
866 }
867
868 /// Lists input columns whose values are used in outputs.
869 ///
870 /// You can use `BTreeSet::last()` to extract the maximum demanded column from the set.
871 ///
872 /// It is entirely appropriate to determine the demand of an instance
873 /// and then both apply a projection to the subject of the instance and
874 /// `self.permute` this instance.
875 pub fn demand(&self) -> BTreeSet<usize> {
876 self.demand_with(BTreeSet::new())
877 }
878
879 /// Like [`Self::demand`], but also treats the columns in `demanded`, which
880 /// index the row of inputs followed by mapped values, as used.
881 fn demand_with(&self, mut demanded: BTreeSet<usize>) -> BTreeSet<usize> {
882 for (_index, pred) in self.predicates.iter() {
883 demanded.extend(pred.support());
884 }
885 demanded.extend(self.projection.iter().cloned());
886 for index in (0..self.expressions.len()).rev() {
887 if demanded.contains(&(self.input_arity + index)) {
888 demanded.extend(self.expressions[index].support());
889 }
890 }
891 demanded.retain(|col| col < &self.input_arity);
892 demanded
893 }
894
895 /// Update input column references, due to an input projection or permutation.
896 ///
897 /// The `shuffle` argument remaps expected column identifiers to new locations,
898 /// with the expectation that `shuffle` describes all input columns, and so the
899 /// intermediate results will be able to start at position `shuffle.len()`.
900 ///
901 /// The supplied `shuffle` might not list columns that are not "demanded" by the
902 /// instance, and so we should ensure that `self` is optimized to not reference
903 /// columns that are not demanded.
904 pub fn permute_fn<F>(&mut self, remap: F, new_input_arity: usize)
905 where
906 F: Fn(usize) -> usize,
907 {
908 let (mut map, mut filter, mut project) = self.as_map_filter_project();
909 let action = self.permute_action(&remap, new_input_arity);
910 for expr in map.iter_mut() {
911 expr.visit_columns(&action);
912 }
913 for pred in filter.iter_mut() {
914 pred.visit_columns(&action);
915 }
916 for proj in project.iter_mut() {
917 action(proj);
918 assert!(*proj < new_input_arity + map.len());
919 }
920 *self = Self::new(new_input_arity)
921 .map(map)
922 .filter(filter)
923 .project(project)
924 }
925
926 /// The column renaming [`Self::permute_fn`] applies to every reference into
927 /// the row of inputs followed by mapped values.
928 fn permute_action<'a, F>(
929 &self,
930 remap: &'a F,
931 new_input_arity: usize,
932 ) -> impl Fn(&mut usize) + use<'a, F, E>
933 where
934 F: Fn(usize) -> usize,
935 {
936 let input_arity = self.input_arity;
937 let map_len = self.expressions.len();
938 move |col: &mut usize| {
939 if input_arity <= *col && *col < input_arity + map_len {
940 *col = new_input_arity + (*col - input_arity);
941 } else {
942 *col = remap(*col);
943 }
944 }
945 }
946}
947
948// Optimization routines.
949impl<E: OptimizableExpr> MapFilterProject<E> {
950 /// Optimize the internal expression evaluation order.
951 ///
952 /// This method performs several optimizations that are meant to streamline
953 /// the execution of the `MapFilterProject` instance, but not to alter its
954 /// semantics. This includes extracting expressions that are used multiple
955 /// times, inlining those that are not, and removing expressions that are
956 /// unreferenced.
957 ///
958 /// This method will inline all temporal expressions, and remove any columns
959 /// that are not demanded by the output, which should transform any temporal
960 /// filters to a state where the temporal expressions exist only in the list
961 /// of predicates.
962 ///
963 /// # Example
964 ///
965 /// This example demonstrates how the re-use of one expression, converting
966 /// column 1 from a string to an integer, can be extracted and the results
967 /// shared among the two uses. This example is used for each of the steps
968 /// along the optimization path.
969 ///
970 /// ```rust
971 /// use mz_expr::{func, MapFilterProject, MirScalarExpr, UnaryFunc, BinaryFunc};
972 /// // Demonstrate extraction of common expressions (here: parsing strings).
973 /// let mut map_filter_project = MapFilterProject::new(5)
974 /// .map(vec![
975 /// MirScalarExpr::column(0).call_unary(func::CastStringToInt64).call_binary(MirScalarExpr::column(1).call_unary(func::CastStringToInt64), func::AddInt64),
976 /// MirScalarExpr::column(1).call_unary(func::CastStringToInt64).call_binary(MirScalarExpr::column(2).call_unary(func::CastStringToInt64), func::AddInt64),
977 /// ])
978 /// .project(vec![3,4,5,6]);
979 ///
980 /// let mut expected_optimized = MapFilterProject::new(5)
981 /// .map(vec![
982 /// MirScalarExpr::column(1).call_unary(func::CastStringToInt64),
983 /// MirScalarExpr::column(0).call_unary(func::CastStringToInt64).call_binary(MirScalarExpr::column(5), func::AddInt64),
984 /// MirScalarExpr::column(5).call_binary(MirScalarExpr::column(2).call_unary(func::CastStringToInt64), func::AddInt64),
985 /// ])
986 /// .project(vec![3,4,6,7]);
987 ///
988 /// // Optimize the expression.
989 /// map_filter_project.optimize();
990 ///
991 /// assert_eq!(
992 /// map_filter_project,
993 /// expected_optimized,
994 /// );
995 /// ```
996 pub fn optimize(&mut self) {
997 // Track sizes and iterate as long as they decrease.
998 let mut prev_size = None;
999 let mut self_size = usize::MAX;
1000 // Continue as long as strict improvements occur.
1001 while prev_size.map(|p| self_size < p).unwrap_or(true) {
1002 // Lock in current size.
1003 prev_size = Some(self_size);
1004
1005 // We have an annoying pattern of mapping literals that already exist as columns (by filters).
1006 // Try to identify this pattern, of a map that introduces an expression equated to a prior column,
1007 // and then replace the mapped expression by a column reference.
1008 //
1009 // We think this is due to `LiteralLifting`, and we might investigate removing the introduciton in
1010 // the first place. The tell-tale that we see when we fix is a diff that look likes
1011 //
1012 // - Project (#0, #2)
1013 // - Filter (#1 = 1)
1014 // - Map (1)
1015 // - Get l0
1016 // + Filter (#1 = 1)
1017 // + Get l0
1018 //
1019 for (index, expr) in self.expressions.iter_mut().enumerate() {
1020 // If `expr` matches a filter equating it to a column < index + input_arity, rewrite it
1021 for (_, predicate) in self.predicates.iter() {
1022 if let Some(col) =
1023 E::equality_column_alias(predicate, expr, index + self.input_arity)
1024 {
1025 *expr = col;
1026 }
1027 }
1028 }
1029
1030 // Optimization memoizes individual `ScalarExpr` expressions that
1031 // are sure to be evaluated, canonicalizes references to the first
1032 // occurrence of each, inlines expressions that have a reference
1033 // count of one, and then removes any expressions that are not
1034 // referenced.
1035 self.memoize_expressions();
1036 self.predicates.sort();
1037 self.predicates.dedup();
1038 self.inline_expressions();
1039 self.remove_undemanded();
1040
1041 // Re-build `self` from parts to restore evaluation order invariants.
1042 let (map, filter, project) = self.as_map_filter_project();
1043 *self = Self::new(self.input_arity)
1044 .map(map)
1045 .filter(filter)
1046 .project(project);
1047
1048 self_size = self.size();
1049 }
1050 }
1051
1052 /// Total expression sizes across all expressions.
1053 pub fn size(&self) -> usize {
1054 self.expressions
1055 .iter()
1056 .map(|e| OptimizableExpr::size(e))
1057 .sum::<usize>()
1058 + self
1059 .predicates
1060 .iter()
1061 .map(|(_, e)| OptimizableExpr::size(e))
1062 .sum::<usize>()
1063 }
1064
1065 /// Place each certainly evaluated expression in its own column.
1066 ///
1067 /// This method places each non-trivial, certainly evaluated expression
1068 /// in its own column, and deduplicates them so that all references to
1069 /// the same expression reference the same column.
1070 ///
1071 /// This transformation is restricted to expressions we are certain will
1072 /// be evaluated, which does not include expressions in `if` statements.
1073 ///
1074 /// # Example
1075 ///
1076 /// This example demonstrates how memoization notices `MirScalarExpr`s
1077 /// that are used multiple times, and ensures that each are extracted
1078 /// into columns and then referenced by column. This pass does not try
1079 /// to minimize the occurrences of column references, which will happen
1080 /// in inlining.
1081 ///
1082 /// ```rust
1083 /// use mz_expr::{func, MapFilterProject, MirScalarExpr, UnaryFunc, BinaryFunc};
1084 /// // Demonstrate extraction of common expressions (here: parsing strings).
1085 /// let mut map_filter_project = MapFilterProject::new(5)
1086 /// .map(vec![
1087 /// MirScalarExpr::column(0).call_unary(func::CastStringToInt64).call_binary(MirScalarExpr::column(1).call_unary(func::CastStringToInt64), func::AddInt64),
1088 /// MirScalarExpr::column(1).call_unary(func::CastStringToInt64).call_binary(MirScalarExpr::column(2).call_unary(func::CastStringToInt64), func::AddInt64),
1089 /// ])
1090 /// .project(vec![3,4,5,6]);
1091 ///
1092 /// let mut expected_optimized = MapFilterProject::new(5)
1093 /// .map(vec![
1094 /// MirScalarExpr::column(0).call_unary(func::CastStringToInt64),
1095 /// MirScalarExpr::column(1).call_unary(func::CastStringToInt64),
1096 /// MirScalarExpr::column(5).call_binary(MirScalarExpr::column(6), func::AddInt64),
1097 /// MirScalarExpr::column(7),
1098 /// MirScalarExpr::column(2).call_unary(func::CastStringToInt64),
1099 /// MirScalarExpr::column(6).call_binary(MirScalarExpr::column(9), func::AddInt64),
1100 /// MirScalarExpr::column(10),
1101 /// ])
1102 /// .project(vec![3,4,8,11]);
1103 ///
1104 /// // Memoize expressions, ensuring uniqueness of each `MirScalarExpr`.
1105 /// map_filter_project.memoize_expressions();
1106 ///
1107 /// assert_eq!(
1108 /// map_filter_project,
1109 /// expected_optimized,
1110 /// );
1111 /// ```
1112 ///
1113 /// Expressions may not be memoized if they are not certain to be evaluated,
1114 /// for example if they occur in conditional branches of a `MirScalarExpr::If`.
1115 ///
1116 /// ```rust
1117 /// use mz_expr::{func, MapFilterProject, MirScalarExpr, UnaryFunc, BinaryFunc};
1118 /// // Demonstrate extraction of unconditionally evaluated expressions, as well as
1119 /// // the non-extraction of common expressions guarded by conditions.
1120 /// let mut map_filter_project = MapFilterProject::new(2)
1121 /// .map(vec![
1122 /// MirScalarExpr::If {
1123 /// cond: Box::new(MirScalarExpr::column(0).call_binary(MirScalarExpr::column(1), func::Lt)),
1124 /// then: Box::new(MirScalarExpr::column(0).call_binary(MirScalarExpr::column(1), func::DivInt64)),
1125 /// els: Box::new(MirScalarExpr::column(1).call_binary(MirScalarExpr::column(0), func::DivInt64)),
1126 /// },
1127 /// MirScalarExpr::If {
1128 /// cond: Box::new(MirScalarExpr::column(0).call_binary(MirScalarExpr::column(1), func::Lt)),
1129 /// then: Box::new(MirScalarExpr::column(1).call_binary(MirScalarExpr::column(0), func::DivInt64)),
1130 /// els: Box::new(MirScalarExpr::column(0).call_binary(MirScalarExpr::column(1), func::DivInt64)),
1131 /// },
1132 /// ]);
1133 ///
1134 /// let mut expected_optimized = MapFilterProject::new(2)
1135 /// .map(vec![
1136 /// MirScalarExpr::column(0).call_binary(MirScalarExpr::column(1), func::Lt),
1137 /// MirScalarExpr::If {
1138 /// cond: Box::new(MirScalarExpr::column(2)),
1139 /// then: Box::new(MirScalarExpr::column(0).call_binary(MirScalarExpr::column(1), func::DivInt64)),
1140 /// els: Box::new(MirScalarExpr::column(1).call_binary(MirScalarExpr::column(0), func::DivInt64)),
1141 /// },
1142 /// MirScalarExpr::column(3),
1143 /// MirScalarExpr::If {
1144 /// cond: Box::new(MirScalarExpr::column(2)),
1145 /// then: Box::new(MirScalarExpr::column(1).call_binary(MirScalarExpr::column(0), func::DivInt64)),
1146 /// els: Box::new(MirScalarExpr::column(0).call_binary(MirScalarExpr::column(1), func::DivInt64)),
1147 /// },
1148 /// MirScalarExpr::column(5),
1149 /// ])
1150 /// .project(vec![0,1,4,6]);
1151 ///
1152 /// // Memoize expressions, ensuring uniqueness of each `MirScalarExpr`.
1153 /// map_filter_project.memoize_expressions();
1154 ///
1155 /// assert_eq!(
1156 /// map_filter_project,
1157 /// expected_optimized,
1158 /// );
1159 /// ```
1160 pub fn memoize_expressions(&mut self) {
1161 // Record the mapping from starting column references to new column
1162 // references.
1163 let mut remaps = BTreeMap::new();
1164 for index in 0..self.input_arity {
1165 remaps.insert(index, index);
1166 }
1167 let mut new_expressions = Vec::new();
1168
1169 // We follow the same order as for evaluation, to ensure that all
1170 // column references exist in time for their evaluation. We could
1171 // prioritize predicates, but we would need to be careful to chase
1172 // down column references to expressions and memoize those as well.
1173 let mut expression = 0;
1174 for (support, predicate) in self.predicates.iter_mut() {
1175 while self.input_arity + expression < *support {
1176 self.expressions[expression].permute_map(&remaps);
1177 memoize_expr(
1178 &mut self.expressions[expression],
1179 &mut new_expressions,
1180 self.input_arity,
1181 );
1182 remaps.insert(
1183 self.input_arity + expression,
1184 self.input_arity + new_expressions.len(),
1185 );
1186 new_expressions.push(self.expressions[expression].clone());
1187 expression += 1;
1188 }
1189 predicate.permute_map(&remaps);
1190 memoize_expr(predicate, &mut new_expressions, self.input_arity);
1191 }
1192 while expression < self.expressions.len() {
1193 self.expressions[expression].permute_map(&remaps);
1194 memoize_expr(
1195 &mut self.expressions[expression],
1196 &mut new_expressions,
1197 self.input_arity,
1198 );
1199 remaps.insert(
1200 self.input_arity + expression,
1201 self.input_arity + new_expressions.len(),
1202 );
1203 new_expressions.push(self.expressions[expression].clone());
1204 expression += 1;
1205 }
1206
1207 self.expressions = new_expressions;
1208 for proj in self.projection.iter_mut() {
1209 *proj = remaps[proj];
1210 }
1211
1212 // Restore predicate order invariants.
1213 for (pos, pred) in self.predicates.iter_mut() {
1214 *pos = pred.support().last().map(|x| *x + 1).unwrap_or(0);
1215 }
1216 }
1217
1218 /// This method inlines expressions with a single use.
1219 ///
1220 /// This method only inlines expressions; it does not delete expressions
1221 /// that are no longer referenced. The `remove_undemanded()` method does
1222 /// that, and should likely be used after this method.
1223 ///
1224 /// Inlining replaces column references when the referred-to item is either
1225 /// another column reference, or the only referrer of its referent. This
1226 /// is most common after memoization has atomized all expressions to seek
1227 /// out re-use: inlining re-assembles expressions that were not helpfully
1228 /// shared with other expressions.
1229 ///
1230 /// # Example
1231 ///
1232 /// In this example, we see that with only a single reference to columns
1233 /// 0 and 2, their parsing can each be inlined. Similarly, column references
1234 /// can be cleaned up among expressions, and in the final projection.
1235 ///
1236 /// Also notice the remaining expressions, which can be cleaned up in a later
1237 /// pass (the `remove_undemanded` method).
1238 ///
1239 /// ```rust
1240 /// use mz_expr::{func, MapFilterProject, MirScalarExpr, UnaryFunc, BinaryFunc};
1241 /// // Use the output from first `memoize_expression` example.
1242 /// let mut map_filter_project = MapFilterProject::new(5)
1243 /// .map(vec![
1244 /// MirScalarExpr::column(0).call_unary(func::CastStringToInt64),
1245 /// MirScalarExpr::column(1).call_unary(func::CastStringToInt64),
1246 /// MirScalarExpr::column(5).call_binary(MirScalarExpr::column(6), func::AddInt64),
1247 /// MirScalarExpr::column(7),
1248 /// MirScalarExpr::column(2).call_unary(func::CastStringToInt64),
1249 /// MirScalarExpr::column(6).call_binary(MirScalarExpr::column(9), func::AddInt64),
1250 /// MirScalarExpr::column(10),
1251 /// ])
1252 /// .project(vec![3,4,8,11]);
1253 ///
1254 /// let mut expected_optimized = MapFilterProject::new(5)
1255 /// .map(vec![
1256 /// MirScalarExpr::column(0).call_unary(func::CastStringToInt64),
1257 /// MirScalarExpr::column(1).call_unary(func::CastStringToInt64),
1258 /// MirScalarExpr::column(0).call_unary(func::CastStringToInt64).call_binary(MirScalarExpr::column(6), func::AddInt64),
1259 /// MirScalarExpr::column(0).call_unary(func::CastStringToInt64).call_binary(MirScalarExpr::column(6), func::AddInt64),
1260 /// MirScalarExpr::column(2).call_unary(func::CastStringToInt64),
1261 /// MirScalarExpr::column(6).call_binary(MirScalarExpr::column(2).call_unary(func::CastStringToInt64), func::AddInt64),
1262 /// MirScalarExpr::column(6).call_binary(MirScalarExpr::column(2).call_unary(func::CastStringToInt64), func::AddInt64),
1263 /// ])
1264 /// .project(vec![3,4,8,11]);
1265 ///
1266 /// // Inline expressions that are referenced only once.
1267 /// map_filter_project.inline_expressions();
1268 ///
1269 /// assert_eq!(
1270 /// map_filter_project,
1271 /// expected_optimized,
1272 /// );
1273 /// ```
1274 pub fn inline_expressions(&mut self) {
1275 // Local copy of input_arity to avoid borrowing `self` in closures.
1276 let input_arity = self.input_arity;
1277 // Reference counts track the number of places that a reference occurs.
1278 let mut reference_count = vec![0; input_arity + self.expressions.len()];
1279 // Increment reference counts for each use
1280 for expr in self.expressions.iter() {
1281 expr.visit_pre(&mut |e| {
1282 if let Some(i) = e.as_column() {
1283 reference_count[i] += 1;
1284 }
1285 });
1286 }
1287 for (_, pred) in self.predicates.iter() {
1288 pred.visit_pre(&mut |e| {
1289 if let Some(i) = e.as_column() {
1290 reference_count[i] += 1;
1291 }
1292 });
1293 }
1294 for proj in self.projection.iter() {
1295 reference_count[*proj] += 1;
1296 }
1297
1298 // Determine which expressions should be inlined because they reference temporal expressions.
1299 let mut is_temporal = vec![false; input_arity];
1300 for expr in self.expressions.iter() {
1301 // An express may contain a temporal expression, or reference a column containing such.
1302 is_temporal.push(
1303 OptimizableExpr::contains_temporal(expr)
1304 || expr.support().into_iter().any(|col| is_temporal[col]),
1305 );
1306 }
1307
1308 // Inline only those columns that 1. are expressions not inputs, and
1309 // 2a. are column references or literals or 2b. have a refcount of 1,
1310 // or 2c. reference temporal expressions (which cannot be evaluated).
1311 let mut should_inline = vec![false; reference_count.len()];
1312 for i in (input_arity..reference_count.len()).rev() {
1313 if let Some(c) = self.expressions[i - input_arity].as_column() {
1314 should_inline[i] = true;
1315 // The reference count of the referenced column should be
1316 // incremented with the number of references
1317 // `self.expressions[i - input_arity]` has.
1318 // Subtract 1 because `self.expressions[i - input_arity]` is
1319 // itself a reference.
1320 reference_count[c] += reference_count[i] - 1;
1321 } else {
1322 should_inline[i] = reference_count[i] == 1 || is_temporal[i];
1323 }
1324 }
1325 // Inline expressions per `should_inline`.
1326 self.perform_inlining(should_inline);
1327 // We can only inline column references in `self.projection`, but we should.
1328 for proj in self.projection.iter_mut() {
1329 if *proj >= self.input_arity {
1330 if let Some(i) = self.expressions[*proj - self.input_arity].as_column() {
1331 // TODO(mgree) !!! propagate name information to projection
1332 *proj = i;
1333 }
1334 }
1335 }
1336 }
1337
1338 /// Inlines those expressions that are indicated by should_inline.
1339 /// See `inline_expressions` for usage.
1340 pub fn perform_inlining(&mut self, should_inline: Vec<bool>) {
1341 for index in 0..self.expressions.len() {
1342 let (prior, expr) = self.expressions.split_at_mut(index);
1343 expr[0].visit_mut_post(&mut |e| {
1344 if let Some(i) = e.as_column() {
1345 if should_inline[i] {
1346 *e = prior[i - self.input_arity].clone();
1347 }
1348 }
1349 });
1350 }
1351 for (_index, pred) in self.predicates.iter_mut() {
1352 let expressions = &self.expressions;
1353 pred.visit_mut_post(&mut |e| {
1354 if let Some(i) = e.as_column() {
1355 if should_inline[i] {
1356 *e = expressions[i - self.input_arity].clone();
1357 }
1358 }
1359 });
1360 }
1361 }
1362
1363 /// Removes unused expressions from `self.expressions`.
1364 ///
1365 /// Expressions are "used" if they are relied upon by any output columns
1366 /// or any predicates, even transitively. Any expressions that are not
1367 /// relied upon in this way can be discarded.
1368 ///
1369 /// # Example
1370 ///
1371 /// ```rust
1372 /// use mz_expr::{func, MapFilterProject, MirScalarExpr, UnaryFunc, BinaryFunc};
1373 /// // Use the output from `inline_expression` example.
1374 /// let mut map_filter_project = MapFilterProject::new(5)
1375 /// .map(vec![
1376 /// MirScalarExpr::column(0).call_unary(func::CastStringToInt64),
1377 /// MirScalarExpr::column(1).call_unary(func::CastStringToInt64),
1378 /// MirScalarExpr::column(0).call_unary(func::CastStringToInt64).call_binary(MirScalarExpr::column(6), func::AddInt64),
1379 /// MirScalarExpr::column(0).call_unary(func::CastStringToInt64).call_binary(MirScalarExpr::column(6), func::AddInt64),
1380 /// MirScalarExpr::column(2).call_unary(func::CastStringToInt64),
1381 /// MirScalarExpr::column(6).call_binary(MirScalarExpr::column(2).call_unary(func::CastStringToInt64), func::AddInt64),
1382 /// MirScalarExpr::column(6).call_binary(MirScalarExpr::column(2).call_unary(func::CastStringToInt64), func::AddInt64),
1383 /// ])
1384 /// .project(vec![3,4,8,11]);
1385 ///
1386 /// let mut expected_optimized = MapFilterProject::new(5)
1387 /// .map(vec![
1388 /// MirScalarExpr::column(1).call_unary(func::CastStringToInt64),
1389 /// MirScalarExpr::column(0).call_unary(func::CastStringToInt64).call_binary(MirScalarExpr::column(5), func::AddInt64),
1390 /// MirScalarExpr::column(5).call_binary(MirScalarExpr::column(2).call_unary(func::CastStringToInt64), func::AddInt64),
1391 /// ])
1392 /// .project(vec![3,4,6,7]);
1393 ///
1394 /// // Remove undemanded expressions, streamlining the work done..
1395 /// map_filter_project.remove_undemanded();
1396 ///
1397 /// assert_eq!(
1398 /// map_filter_project,
1399 /// expected_optimized,
1400 /// );
1401 /// ```
1402 pub fn remove_undemanded(&mut self) {
1403 // Determine the demanded expressions to remove irrelevant ones.
1404 let mut demand = BTreeSet::new();
1405 for (_index, pred) in self.predicates.iter() {
1406 demand.extend(pred.support());
1407 }
1408 // Start from the output columns as presumed demanded.
1409 // If this is not the case, the caller should project some away.
1410 demand.extend(self.projection.iter().cloned());
1411 // Proceed in *reverse* order, as expressions may depend on other
1412 // expressions that precede them.
1413 for index in (0..self.expressions.len()).rev() {
1414 if demand.contains(&(self.input_arity + index)) {
1415 demand.extend(self.expressions[index].support());
1416 }
1417 }
1418
1419 // Maintain a map from initial column identifiers to locations
1420 // once we have removed undemanded expressions.
1421 let mut remap = BTreeMap::new();
1422 // This map only needs to map elements of `demand` to a new location,
1423 // but the logic is easier if we include all input columns (as the
1424 // new position is then determined by the size of the map).
1425 for index in 0..self.input_arity {
1426 remap.insert(index, index);
1427 }
1428 // Retain demanded expressions, and record their new locations.
1429 let mut new_expressions = Vec::new();
1430 for (index, expr) in self.expressions.drain(..).enumerate() {
1431 if demand.contains(&(index + self.input_arity)) {
1432 remap.insert(index + self.input_arity, remap.len());
1433 new_expressions.push(expr);
1434 }
1435 }
1436 self.expressions = new_expressions;
1437
1438 // Update column identifiers; rebuild `Self` to re-establish any invariants.
1439 // We mirror `self.permute(&remap)` but we specifically want to remap columns
1440 // that are produced by `self.expressions` after the input columns.
1441 let (expressions, predicates, projection) = self.as_map_filter_project();
1442 *self = Self::new(self.input_arity)
1443 .map(expressions.into_iter().map(|mut e| {
1444 e.permute_map(&remap);
1445 e
1446 }))
1447 .filter(predicates.into_iter().map(|mut p| {
1448 p.permute_map(&remap);
1449 p
1450 }))
1451 .project(projection.into_iter().map(|c| remap[&c]));
1452 }
1453}
1454
1455// TODO: move this elsewhere?
1456/// Recursively memoize parts of `expr`, storing those parts in `memoized_parts`.
1457///
1458/// A part of `expr` that is memoized is replaced by a reference to column
1459/// `(input_arity + pos)`, where `pos` is the position of the memoized part in
1460/// `memoized_parts`, and `input_arity` is the arity of the input that `expr`
1461/// refers to.
1462pub fn memoize_expr<E: OptimizableExpr>(
1463 expr: &mut E,
1464 memoized_parts: &mut Vec<E>,
1465 input_arity: usize,
1466) {
1467 expr.visit_mut_pre_post(&mut |e| e.eager_children(), &mut |e| {
1468 if E::is_literal(e) {
1469 // Literals do not need to be memoized.
1470 return;
1471 }
1472 if let Some(col) = e.as_column_mut() {
1473 // Column references do not need to be memoized, but may need to be
1474 // updated if they reference a column reference themselves.
1475 if *col > input_arity {
1476 if let Some(col2) = memoized_parts[*col - input_arity].as_column() {
1477 // Update the column index in place, preserving any name information.
1478 *col = col2;
1479 }
1480 }
1481 return;
1482 }
1483 // TODO: OOO (Optimizer Optimization Opportunity):
1484 // we are quadratic in expression size because of this .iter().position
1485 if let Some(position) = memoized_parts.iter().position(|e2| e2 == e) {
1486 // Any complex expression that already exists as a prior column can
1487 // be replaced by a reference to that column.
1488 *e = E::column(input_arity + position);
1489 } else {
1490 // A complex expression that does not exist should be memoized, and
1491 // replaced by a reference to the column.
1492 memoized_parts.push(std::mem::replace(
1493 e,
1494 E::column(input_arity + memoized_parts.len()),
1495 ));
1496 }
1497 });
1498}
1499
1500pub mod util {
1501 use std::collections::BTreeMap;
1502
1503 use crate::MirScalarExpr;
1504 use crate::scalar::columns::Columns;
1505
1506 #[allow(dead_code)]
1507 /// A triple of actions that map from rows to (key, val) pairs and back again.
1508 struct KeyValRowMapping {
1509 /// Expressions to apply to a row to produce key datums.
1510 to_key: Vec<MirScalarExpr>,
1511 /// Columns to project from a row to produce residual value datums.
1512 to_val: Vec<usize>,
1513 /// Columns to project from the concatenation of key and value to reconstruct the row.
1514 to_row: Vec<usize>,
1515 }
1516
1517 /// Derive supporting logic to support transforming rows to (key, val) pairs,
1518 /// and back again.
1519 ///
1520 /// We are given as input a list of mappings from columns to key indices, a key length,
1521 /// and an input arity. (The produced key should be the application of the key expressions.)
1522 ///
1523 /// To produce the `val` output, we will identify those input columns not found in
1524 /// the key expressions, and name all other columns.
1525 /// To reconstitute the original row, we identify the sequence of columns from the
1526 /// concatenation of key and val which would reconstruct the original row.
1527 ///
1528 /// The output is a pair of column sequences, the first used to reconstruct a row
1529 /// from the concatenation of key and value, and the second to identify the columns
1530 /// of a row that should become the value associated with its key.
1531 ///
1532 /// The permutations and thinning expressions generated here will be tracked in
1533 /// `compute_types::plan::AvailableCollections`; see the
1534 /// documentation there for more details.
1535 pub fn permutation_for_arrangement(
1536 key: &[impl Columns],
1537 unthinned_arity: usize,
1538 ) -> (Vec<usize>, Vec<usize>) {
1539 let columns_in_key: BTreeMap<_, _> = key
1540 .iter()
1541 .enumerate()
1542 .filter_map(|(i, key_col)| key_col.as_column().map(|c| (c, i)))
1543 .collect();
1544
1545 let mut input_cursor = key.len();
1546 let permutation = (0..unthinned_arity)
1547 .map(|c| {
1548 if let Some(c) = columns_in_key.get(&c) {
1549 // Column is in key (and thus gone from the value
1550 // of the thinned representation)
1551 *c
1552 } else {
1553 // Column remains in value of the thinned representation
1554 input_cursor += 1;
1555 input_cursor - 1
1556 }
1557 })
1558 .collect();
1559 let thinning = (0..unthinned_arity)
1560 .filter(|c| !columns_in_key.contains_key(c))
1561 .collect();
1562 (permutation, thinning)
1563 }
1564
1565 /// Given the permutations (see [`permutation_for_arrangement`] and
1566 /// (`dataflow::plan::AvailableCollections`) corresponding to two
1567 /// collections with the same key arity,
1568 /// computes the permutation for the result of joining them.
1569 pub fn join_permutations(
1570 key_arity: usize,
1571 stream_permutation: Vec<usize>,
1572 thinned_stream_arity: usize,
1573 lookup_permutation: Vec<usize>,
1574 ) -> BTreeMap<usize, usize> {
1575 let stream_arity = stream_permutation.len();
1576 let lookup_arity = lookup_permutation.len();
1577
1578 (0..stream_arity + lookup_arity)
1579 .map(|i| {
1580 let location = if i < stream_arity {
1581 stream_permutation[i]
1582 } else {
1583 let location_in_lookup = lookup_permutation[i - stream_arity];
1584 if location_in_lookup < key_arity {
1585 location_in_lookup
1586 } else {
1587 location_in_lookup + thinned_stream_arity
1588 }
1589 };
1590 (i, location)
1591 })
1592 .collect()
1593 }
1594}
1595
1596pub mod plan {
1597 use std::collections::BTreeSet;
1598 use std::iter;
1599
1600 use mz_repr::{Datum, Diff, Row, RowArena};
1601 use serde::{Deserialize, Serialize};
1602
1603 use crate::scalar::optimizable::OptimizableExpr;
1604 use crate::{Columns, Eval};
1605 use crate::{EvalError, MapFilterProject, MirScalarExpr};
1606
1607 /// A wrapper type which indicates it is safe to simply evaluate all expressions.
1608 #[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq, Ord, PartialOrd)]
1609 #[serde(bound(deserialize = "E: serde::de::DeserializeOwned"))]
1610 pub struct SafeMfpPlan<E: OptimizableExpr = MirScalarExpr> {
1611 pub(crate) mfp: MapFilterProject<E>,
1612 }
1613
1614 impl<E: OptimizableExpr> SafeMfpPlan<E> {
1615 /// Wrap a `MapFilterProject` in a `SafeMfpPlan`.
1616 pub fn from_mfp(mfp: MapFilterProject<E>) -> Self {
1617 Self { mfp }
1618 }
1619
1620 /// Unwrap the inner `MapFilterProject`.
1621 pub fn into_mfp(self) -> MapFilterProject<E> {
1622 self.mfp
1623 }
1624
1625 /// Remaps references to input columns according to `remap`.
1626 ///
1627 /// Leaves other column references, e.g. to newly mapped columns, unchanged.
1628 pub fn permute_fn<F>(&mut self, remap: F, new_arity: usize)
1629 where
1630 F: Fn(usize) -> usize,
1631 {
1632 self.mfp.permute_fn(remap, new_arity);
1633 }
1634
1635 /// Returns true when `Self` is the identity.
1636 pub fn is_identity(&self) -> bool {
1637 self.mfp.is_identity()
1638 }
1639 }
1640
1641 impl<E: OptimizableExpr + Eval> SafeMfpPlan<E> {
1642 /// Evaluates the linear operator on a supplied list of datums.
1643 ///
1644 /// The arguments are the initial datums associated with the row,
1645 /// and an appropriately lifetimed arena for temporary allocations
1646 /// needed by scalar evaluation.
1647 ///
1648 /// An `Ok` result will either be `None` if any predicate did not
1649 /// evaluate to `Datum::True`, or the values of the columns listed
1650 /// by `self.projection` if all predicates passed. If an error
1651 /// occurs in the evaluation it is returned as an `Err` variant.
1652 /// As the evaluation exits early with failed predicates, it may
1653 /// miss some errors that would occur later in evaluation.
1654 ///
1655 /// The `row` is not cleared first, but emptied if the function
1656 /// returns `Ok(Some(row)).
1657 #[inline(always)]
1658 pub fn evaluate_into<'a, 'row>(
1659 &'a self,
1660 datums: &mut Vec<Datum<'a>>,
1661 arena: &'a RowArena,
1662 row_buf: &'row mut Row,
1663 ) -> Result<Option<&'row Row>, EvalError> {
1664 let passed_predicates = self.evaluate_inner(datums, arena)?;
1665 if !passed_predicates {
1666 Ok(None)
1667 } else {
1668 row_buf
1669 .packer()
1670 .extend(self.mfp.projection.iter().map(|c| datums[*c]));
1671 Ok(Some(row_buf))
1672 }
1673 }
1674
1675 /// A version of `evaluate` which produces an iterator over `Datum`
1676 /// as output.
1677 ///
1678 /// This version can be useful when one wants to capture the resulting
1679 /// datums without packing and then unpacking a row.
1680 #[inline(always)]
1681 pub fn evaluate_iter<'b, 'a: 'b>(
1682 &'a self,
1683 datums: &'b mut Vec<Datum<'a>>,
1684 arena: &'a RowArena,
1685 ) -> Result<Option<impl Iterator<Item = Datum<'a>> + 'b>, EvalError> {
1686 let passed_predicates = self.evaluate_inner(datums, arena)?;
1687 if !passed_predicates {
1688 Ok(None)
1689 } else {
1690 Ok(Some(self.mfp.projection.iter().map(move |i| datums[*i])))
1691 }
1692 }
1693
1694 /// Populates `datums` with `self.expressions` and tests `self.predicates`.
1695 ///
1696 /// This does not apply `self.projection`, which is up to the calling method.
1697 pub fn evaluate_inner<'b, 'a: 'b>(
1698 &'a self,
1699 datums: &'b mut Vec<Datum<'a>>,
1700 arena: &'a RowArena,
1701 ) -> Result<bool, EvalError> {
1702 let mut expression = 0;
1703 for (support, predicate) in self.mfp.predicates.iter() {
1704 while self.mfp.input_arity + expression < *support {
1705 datums.push(self.mfp.expressions[expression].eval(&datums[..], arena)?);
1706 expression += 1;
1707 }
1708 if predicate.eval(&datums[..], arena)? != Datum::True {
1709 return Ok(false);
1710 }
1711 }
1712 while expression < self.mfp.expressions.len() {
1713 datums.push(self.mfp.expressions[expression].eval(&datums[..], arena)?);
1714 expression += 1;
1715 }
1716 Ok(true)
1717 }
1718
1719 /// Returns true if evaluation could introduce an error on non-error inputs.
1720 pub fn could_error(&self) -> bool {
1721 self.mfp.predicates.iter().any(|(_pos, e)| e.could_error())
1722 || self.mfp.expressions.iter().any(|e| e.could_error())
1723 }
1724 }
1725
1726 impl<E: OptimizableExpr> std::ops::Deref for SafeMfpPlan<E> {
1727 type Target = MapFilterProject<E>;
1728 fn deref(&self) -> &Self::Target {
1729 &self.mfp
1730 }
1731 }
1732
1733 /// Predicates partitioned into temporal and non-temporal.
1734 ///
1735 /// Temporal predicates require some recognition to determine their
1736 /// structure, and it is best to do that once and re-use the results.
1737 ///
1738 /// There are restrictions on the temporal predicates we currently support.
1739 /// They must directly constrain `MzNow` from below or above,
1740 /// by expressions that do not themselves contain `MzNow`.
1741 /// Conjunctions of such constraints are also ok.
1742 #[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)]
1743 #[serde(bound(deserialize = "E: serde::de::DeserializeOwned"))]
1744 pub struct MfpPlan<E: OptimizableExpr = MirScalarExpr> {
1745 /// Normal predicates to evaluate on `&[Datum]` and expect `Ok(Datum::True)`.
1746 pub(crate) mfp: SafeMfpPlan<E>,
1747 /// Expressions that when evaluated lower-bound or equal (<=) `MzNow`.
1748 pub(crate) lower_bounds: Vec<E>,
1749 /// Expressions that when evaluated strictly upper-bound `MzNow`.
1750 pub(crate) upper_bounds: Vec<E>,
1751 }
1752
1753 impl<E: OptimizableExpr> MfpPlan<E> {
1754 /// Construct an `MfpPlan` from its components.
1755 pub fn from_parts(mfp: SafeMfpPlan<E>, lower_bounds: Vec<E>, upper_bounds: Vec<E>) -> Self {
1756 Self {
1757 mfp,
1758 lower_bounds,
1759 upper_bounds,
1760 }
1761 }
1762
1763 /// Deconstruct into components.
1764 pub fn into_parts(self) -> (SafeMfpPlan<E>, Vec<E>, Vec<E>) {
1765 (self.mfp, self.lower_bounds, self.upper_bounds)
1766 }
1767
1768 /// Access the inner `SafeMfpPlan`.
1769 pub fn safe_mfp(&self) -> &SafeMfpPlan<E> {
1770 &self.mfp
1771 }
1772
1773 /// Borrow all parts: `(safe_mfp, lower_bounds, upper_bounds)`.
1774 pub fn as_parts(&self) -> (&SafeMfpPlan<E>, &[E], &[E]) {
1775 (&self.mfp, &self.lower_bounds, &self.upper_bounds)
1776 }
1777
1778 /// Partitions `predicates` into non-temporal, and lower and upper temporal bounds.
1779 ///
1780 /// The first returned list is of predicates that do not contain `mz_now`.
1781 /// The second and third returned lists contain expressions that, once evaluated, lower
1782 /// and upper bound the validity interval of a record, respectively. These second two
1783 /// lists are populated only by binary expressions of the form
1784 /// ```ignore
1785 /// mz_now cmp_op expr
1786 /// ```
1787 /// where `cmp_op` is a comparison operator and `expr` does not contain `mz_now`.
1788 ///
1789 /// If any unsupported expression is found, for example one that uses `mz_now`
1790 /// in an unsupported position, an error is returned.
1791 pub fn create_from(mut mfp: MapFilterProject<E>) -> Result<Self, String> {
1792 mfp.optimize();
1793
1794 let mut temporal = Vec::new();
1795 mfp.predicates.retain(|(_position, predicate)| {
1796 if OptimizableExpr::contains_temporal(predicate) {
1797 temporal.push(predicate.clone());
1798 false
1799 } else {
1800 true
1801 }
1802 });
1803
1804 let (lower_bounds, upper_bounds) = E::extract_temporal_bounds(temporal)?;
1805
1806 Ok(Self {
1807 mfp: SafeMfpPlan { mfp },
1808 lower_bounds,
1809 upper_bounds,
1810 })
1811 }
1812
1813 /// Lists input columns whose values are used in outputs or temporal bounds.
1814 ///
1815 /// As with [`MapFilterProject::demand`], callers may project their input
1816 /// to these columns and then [`Self::permute_fn`] the plan to match.
1817 pub fn demand(&self) -> BTreeSet<usize> {
1818 let bounds = self.lower_bounds.iter().chain(&self.upper_bounds);
1819 self.mfp
1820 .mfp
1821 .demand_with(bounds.flat_map(|e| e.support()).collect())
1822 }
1823
1824 /// Update input column references, due to an input projection or permutation.
1825 ///
1826 /// See [`MapFilterProject::permute_fn`] for the contract on `remap`.
1827 pub fn permute_fn<F>(&mut self, remap: F, new_arity: usize)
1828 where
1829 F: Fn(usize) -> usize,
1830 {
1831 // Bounds index the same row of inputs followed by mapped values as
1832 // the plan's own expressions, so they take the same renaming. It
1833 // must be computed before the inner plan is permuted.
1834 let action = self.mfp.mfp.permute_action(&remap, new_arity);
1835 for bound in self.lower_bounds.iter_mut().chain(&mut self.upper_bounds) {
1836 bound.visit_columns(&action);
1837 }
1838 self.mfp.permute_fn(&remap, new_arity);
1839 }
1840
1841 /// Indicates if the planned `MapFilterProject` emits exactly its inputs as outputs.
1842 pub fn is_identity(&self) -> bool {
1843 self.mfp.mfp.is_identity()
1844 && self.lower_bounds.is_empty()
1845 && self.upper_bounds.is_empty()
1846 }
1847
1848 /// Returns `true` if the plan contains temporal bounds
1849 /// (i.e., predicates involving `mz_now()`).
1850 pub fn has_temporal_bounds(&self) -> bool {
1851 !self.lower_bounds.is_empty() || !self.upper_bounds.is_empty()
1852 }
1853
1854 /// Returns `self`, and leaves behind an identity operator that acts on its output.
1855 pub fn take(&mut self) -> Self {
1856 let mut identity = Self {
1857 mfp: SafeMfpPlan {
1858 mfp: MapFilterProject::new(self.mfp.projection.len()),
1859 },
1860 lower_bounds: Default::default(),
1861 upper_bounds: Default::default(),
1862 };
1863 std::mem::swap(self, &mut identity);
1864 identity
1865 }
1866
1867 /// Attempt to convert self into a non-temporal MapFilterProject plan.
1868 ///
1869 /// If that is not possible, the original instance is returned as an error.
1870 #[allow(clippy::result_large_err)]
1871 pub fn into_nontemporal(self) -> Result<SafeMfpPlan<E>, Self> {
1872 if self.lower_bounds.is_empty() && self.upper_bounds.is_empty() {
1873 Ok(self.mfp)
1874 } else {
1875 Err(self)
1876 }
1877 }
1878
1879 /// Returns an iterator over mutable references to all non-temporal
1880 /// scalar expressions in the plan.
1881 ///
1882 /// The order of iteration is unspecified.
1883 pub fn iter_nontemporal_exprs(&mut self) -> impl Iterator<Item = &mut E> {
1884 iter::empty()
1885 .chain(self.mfp.mfp.predicates.iter_mut().map(|(_, expr)| expr))
1886 .chain(&mut self.mfp.mfp.expressions)
1887 .chain(&mut self.lower_bounds)
1888 .chain(&mut self.upper_bounds)
1889 }
1890
1891 /// Indicates that `Self` ignores its input to the extent that it can be evaluated on `&[]`.
1892 ///
1893 /// At the moment, this is only true if it projects away all columns and applies no filters,
1894 /// but it could be extended to plans that produce literals independent of the input.
1895 pub fn ignores_input(&self) -> bool {
1896 self.lower_bounds.is_empty()
1897 && self.upper_bounds.is_empty()
1898 && self.mfp.mfp.projection.is_empty()
1899 && self.mfp.mfp.predicates.is_empty()
1900 }
1901 }
1902
1903 impl MfpPlan<MirScalarExpr> {
1904 /// Reconstruct a `MapFilterProject` by folding temporal bounds back as
1905 /// `mz_now()` predicates.
1906 ///
1907 /// This is the inverse of `create_from`: the returned MFP, when passed
1908 /// through `create_from`, will produce an equivalent `MfpPlan`.
1909 ///
1910 /// Lower bounds become `mz_now() >= expr` predicates;
1911 /// upper bounds become `mz_now() < expr` predicates.
1912 pub fn into_map_filter_project(self) -> MapFilterProject<MirScalarExpr> {
1913 let (safe, lower_bounds, upper_bounds) = self.into_parts();
1914 let mut mfp = safe.into_mfp();
1915
1916 let mz_now = MirScalarExpr::CallUnmaterializable(crate::UnmaterializableFunc::MzNow);
1917
1918 for lb in lower_bounds {
1919 // mz_now() >= lb
1920 let predicate = mz_now.clone().call_binary(lb, crate::func::Gte);
1921 let support = predicate
1922 .support()
1923 .into_iter()
1924 .max()
1925 .map(|c| c + 1)
1926 .unwrap_or(0);
1927 mfp.predicates.push((support, predicate));
1928 }
1929 for ub in upper_bounds {
1930 // mz_now() < ub
1931 let predicate = mz_now.clone().call_binary(ub, crate::func::Lt);
1932 let support = predicate
1933 .support()
1934 .into_iter()
1935 .max()
1936 .map(|c| c + 1)
1937 .unwrap_or(0);
1938 mfp.predicates.push((support, predicate));
1939 }
1940
1941 mfp
1942 }
1943 }
1944
1945 impl<E: OptimizableExpr + Eval> MfpPlan<E> {
1946 /// Evaluate the predicates, temporal and non-, and return times and differences for `data`.
1947 ///
1948 /// If `self` contains only non-temporal predicates, the result will either be `(time, diff)`,
1949 /// or an evaluation error. If `self contains temporal predicates, the results can be times
1950 /// that are greater than the input `time`, and may contain negated `diff` values.
1951 ///
1952 /// The `row_builder` is not cleared first, but emptied if the function
1953 /// returns an iterator with any `Ok(_)` element.
1954 pub fn evaluate<'b, 'a: 'b, Err: From<EvalError>, V: Fn(&mz_repr::Timestamp) -> bool>(
1955 &'a self,
1956 datums: &'b mut Vec<Datum<'a>>,
1957 arena: &'a RowArena,
1958 time: mz_repr::Timestamp,
1959 diff: Diff,
1960 valid_time: V,
1961 row_builder: &mut Row,
1962 ) -> impl Iterator<
1963 Item = Result<(Row, mz_repr::Timestamp, Diff), (Err, mz_repr::Timestamp, Diff)>,
1964 > + use<Err, V, E> {
1965 match self.mfp.evaluate_inner(datums, arena) {
1966 Err(e) => {
1967 return Some(Err((e.into(), time, diff))).into_iter().chain(None);
1968 }
1969 Ok(true) => {}
1970 Ok(false) => {
1971 return None.into_iter().chain(None);
1972 }
1973 }
1974 evaluate_temporal(
1975 &self.lower_bounds,
1976 &self.upper_bounds,
1977 &self.mfp.mfp.projection,
1978 datums,
1979 arena,
1980 time,
1981 diff,
1982 valid_time,
1983 row_builder,
1984 )
1985 }
1986
1987 /// Returns true if evaluation could introduce an error on non-error inputs.
1988 pub fn could_error(&self) -> bool {
1989 self.mfp.could_error()
1990 || self.lower_bounds.iter().any(|e| e.could_error())
1991 || self.upper_bounds.iter().any(|e| e.could_error())
1992 }
1993 }
1994
1995 /// Applies temporal bounds to a row whose non-temporal predicates passed,
1996 /// and projects it.
1997 ///
1998 /// `datums` holds every column the bounds and `projection` refer to.
1999 pub(crate) fn evaluate_temporal<
2000 E: Eval,
2001 Err: From<EvalError>,
2002 V: Fn(&mz_repr::Timestamp) -> bool,
2003 >(
2004 lower_bounds: &[E],
2005 upper_bounds: &[E],
2006 projection: &[usize],
2007 datums: &[Datum],
2008 arena: &RowArena,
2009 time: mz_repr::Timestamp,
2010 diff: Diff,
2011 valid_time: V,
2012 row_builder: &mut Row,
2013 ) -> iter::Chain<
2014 std::option::IntoIter<
2015 Result<(Row, mz_repr::Timestamp, Diff), (Err, mz_repr::Timestamp, Diff)>,
2016 >,
2017 std::option::IntoIter<
2018 Result<(Row, mz_repr::Timestamp, Diff), (Err, mz_repr::Timestamp, Diff)>,
2019 >,
2020 > {
2021 // Lower and upper bounds.
2022 let mut lower_bound = time;
2023 let mut upper_bound = None;
2024
2025 // Track whether we have seen a null in either bound, as this should
2026 // prevent the record from being produced at any time.
2027 let mut null_eval = false;
2028
2029 // Advance our lower bound to be at least the result of any lower bound
2030 // expressions.
2031 for l in lower_bounds.iter() {
2032 match l.eval(datums, arena) {
2033 Err(e) => {
2034 return Some(Err((e.into(), time, diff))).into_iter().chain(None);
2035 }
2036 Ok(Datum::MzTimestamp(d)) => {
2037 lower_bound = lower_bound.max(d);
2038 }
2039 Ok(Datum::Null) => {
2040 null_eval = true;
2041 }
2042 x => {
2043 panic!("Non-mz_timestamp value in temporal predicate: {:?}", x);
2044 }
2045 }
2046 }
2047
2048 // If the lower bound exceeds our `until` frontier, it should not appear in the output.
2049 if !valid_time(&lower_bound) {
2050 return None.into_iter().chain(None);
2051 }
2052
2053 // If there are any upper bounds, determine the minimum upper bound.
2054 for u in upper_bounds.iter() {
2055 // We can cease as soon as the lower and upper bounds match,
2056 // as the update will certainly not be produced in that case.
2057 if upper_bound != Some(lower_bound) {
2058 match u.eval(datums, arena) {
2059 Err(e) => {
2060 return Some(Err((e.into(), time, diff))).into_iter().chain(None);
2061 }
2062 Ok(Datum::MzTimestamp(d)) => {
2063 if let Some(upper) = upper_bound {
2064 upper_bound = Some(upper.min(d));
2065 } else {
2066 upper_bound = Some(d);
2067 };
2068 // Force the upper bound to be at least the lower
2069 // bound. The `is_some()` test should always be true
2070 // due to the above block, but maintain it here in
2071 // case that changes. It's hopefully optimized away.
2072 if upper_bound.is_some() && upper_bound < Some(lower_bound) {
2073 upper_bound = Some(lower_bound);
2074 }
2075 }
2076 Ok(Datum::Null) => {
2077 null_eval = true;
2078 }
2079 x => {
2080 panic!("Non-mz_timestamp value in temporal predicate: {:?}", x);
2081 }
2082 }
2083 }
2084 }
2085
2086 // If the upper bound exceeds our `until` frontier, it should not appear in the output.
2087 if let Some(upper) = &mut upper_bound {
2088 if !valid_time(upper) {
2089 upper_bound = None;
2090 }
2091 }
2092
2093 // Produce an output only if the upper bound exceeds the lower bound,
2094 // and if we did not encounter a `null` in our evaluation.
2095 if Some(lower_bound) != upper_bound && !null_eval {
2096 row_builder
2097 .packer()
2098 .extend(projection.iter().map(|c| datums[*c]));
2099 let upper_opt =
2100 upper_bound.map(|upper_bound| Ok((row_builder.clone(), upper_bound, -diff)));
2101 let lower = Some(Ok((row_builder.clone(), lower_bound, diff)));
2102 lower.into_iter().chain(upper_opt)
2103 } else {
2104 None.into_iter().chain(None)
2105 }
2106 }
2107}
2108
2109#[cfg(test)]
2110mod tests {
2111 use mz_repr::{Datum, Diff, Row, RowArena, Timestamp};
2112
2113 use crate::func;
2114 use crate::{EvalError, MapFilterProject, MfpPlan, MirScalarExpr, UnmaterializableFunc};
2115
2116 fn ts(t: u64) -> Datum<'static> {
2117 Datum::MzTimestamp(Timestamp::from(t))
2118 }
2119
2120 /// Reads `ts | int | ts | int | ts | int`, maps `#6 = step_mz_timestamp(#4)`,
2121 /// and keeps `(#3, #6)` while `#0 <= mz_now()`, `mz_now() < step_mz_timestamp(#2)`
2122 /// and `mz_now() < #6`. Columns 0 and 2 are referenced only by bounds, the
2123 /// map column by a bound and the projection, and columns 1 and 5 are unused.
2124 fn temporal_plan() -> MfpPlan {
2125 let mz_now = MirScalarExpr::CallUnmaterializable(UnmaterializableFunc::MzNow);
2126 let step = |c| MirScalarExpr::column(c).call_unary(func::StepMzTimestamp);
2127 let mfp = MapFilterProject::new(6)
2128 .map([step(4)])
2129 .filter([
2130 mz_now
2131 .clone()
2132 .call_binary(MirScalarExpr::column(0), func::Gte),
2133 mz_now.clone().call_binary(step(2), func::Lt),
2134 mz_now.call_binary(MirScalarExpr::column(6), func::Lt),
2135 MirScalarExpr::column(3).call_is_null().not(),
2136 ])
2137 .project([3, 6]);
2138 let plan = MfpPlan::create_from(mfp).expect("plannable");
2139 let (safe, lower, upper) = plan.as_parts();
2140 assert_eq!(safe.expressions.len(), 1, "the map survives planning");
2141 assert_eq!(lower, [MirScalarExpr::column(0)]);
2142 assert!(upper.contains(&MirScalarExpr::column(6)), "{upper:?}");
2143 plan
2144 }
2145
2146 fn rows() -> Vec<Row> {
2147 let int = |i| Datum::Int64(i);
2148 vec![
2149 Row::pack_slice(&[ts(3), int(10), ts(7), int(1), ts(9), int(20)]),
2150 Row::pack_slice(&[ts(3), int(10), ts(9), int(1), ts(7), int(20)]),
2151 Row::pack_slice(&[ts(0), int(11), ts(0), int(2), ts(4), int(21)]),
2152 Row::pack_slice(&[ts(5), int(12), ts(2), int(3), ts(8), int(22)]),
2153 Row::pack_slice(&[ts(1), int(13), ts(9), Datum::Null, ts(9), int(23)]),
2154 ]
2155 }
2156
2157 fn evaluate(
2158 plan: &MfpPlan,
2159 row: &[Datum],
2160 ) -> Vec<Result<(Row, Timestamp, Diff), (EvalError, Timestamp, Diff)>> {
2161 let arena = RowArena::new();
2162 let mut datums = row.to_vec();
2163 plan.evaluate(
2164 &mut datums,
2165 &arena,
2166 Timestamp::from(0u64),
2167 Diff::ONE,
2168 |_| true,
2169 &mut Row::default(),
2170 )
2171 .collect()
2172 }
2173
2174 #[mz_ore::test]
2175 fn mfp_plan_demand_includes_bounds() {
2176 let plan = temporal_plan();
2177 assert_eq!(plan.demand().into_iter().collect::<Vec<_>>(), [0, 2, 3, 4]);
2178 }
2179
2180 #[mz_ore::test]
2181 fn mfp_plan_permute_to_demanded_columns() {
2182 let plan = temporal_plan();
2183 let demand: Vec<usize> = plan.demand().into_iter().collect();
2184 let mut permuted = plan.clone();
2185 let position = |c: usize| demand.iter().position(|d| *d == c).expect("demanded");
2186 permuted.permute_fn(position, demand.len());
2187
2188 for row in rows() {
2189 let full: Vec<Datum> = row.iter().collect();
2190 let narrow: Vec<Datum> = demand.iter().map(|c| full[*c]).collect();
2191 assert_eq!(
2192 evaluate(&permuted, &narrow),
2193 evaluate(&plan, &full),
2194 "{row:?}"
2195 );
2196 }
2197 }
2198
2199 #[mz_ore::test]
2200 fn mfp_plan_truncate_to_max_demand() {
2201 let plan = temporal_plan();
2202 let max_demand = plan.demand().last().map(|c| c + 1).unwrap_or(0);
2203 assert_eq!(max_demand, 5);
2204 let mut truncated = plan.clone();
2205 truncated.permute_fn(|c| c, max_demand);
2206
2207 for row in rows() {
2208 let full: Vec<Datum> = row.iter().collect();
2209 assert_eq!(
2210 evaluate(&truncated, &full[..max_demand]),
2211 evaluate(&plan, &full),
2212 "{row:?}"
2213 );
2214 }
2215 }
2216}