mz_compute_types/plan/lowering.rs
1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Lowering [`DataflowDescription`]s from MIR ([`MirRelationExpr`]) to LIR ([`LirRelationExpr`]).
11
12use std::collections::{BTreeMap, BTreeSet};
13
14use columnar::Len;
15use itertools::Itertools;
16use mz_expr::JoinImplementation::{DeltaQuery, Differential, IndexedFilter, Unimplemented};
17use mz_expr::{
18 AggregateExpr, Columns, Id, JoinInputMapper, MapFilterProject, MfpPlan, MirRelationExpr,
19 MirScalarExpr, OptimizedMirRelationExpr, SafeMfpPlan, TableFunc, permutation_for_arrangement,
20};
21use mz_ore::{assert_none, soft_assert_eq_or_log, soft_panic_or_log};
22use mz_repr::optimize::OptimizerFeatures;
23use mz_repr::{GlobalId, StableRow, Timestamp};
24
25use crate::dataflows::{BuildDesc, DataflowDescription, IndexImport};
26use crate::plan::join::{DeltaJoinPlan, JoinPlan, LinearJoinPlan};
27use crate::plan::reduce::{KeyValPlan, ReducePlan};
28use crate::plan::scalar::{
29 LirScalarExpr, lses_from_mses, mfp_mir_to_lir, mfp_mir_to_lir_plan, mfp_plan_mir_to_lir,
30};
31use crate::plan::threshold::ThresholdPlan;
32use crate::plan::top_k::TopKPlan;
33use crate::plan::{
34 ArrangementStrategy, AvailableCollections, GetPlan, LirId, LirRelationExpr, LirRelationNode,
35 LoweringMetrics,
36};
37
38/// Pick an [`ArrangementStrategy`] based on whether the input may contain future-stamped
39/// updates. Future updates are the only case where temporal bucketing pays off.
40///
41/// Any arrangement or consolidation that absorbs data that can have future updates should be
42/// guarded by a temporal bucketing operator.
43fn strategy_from_future(has_future_updates: bool) -> ArrangementStrategy {
44 if has_future_updates {
45 ArrangementStrategy::TemporalBucketing
46 } else {
47 ArrangementStrategy::Direct
48 }
49}
50
51/// The result of lowering a [`MirRelationExpr`] to a [`LirRelationExpr`].
52struct LoweredExpr {
53 /// The lowered plan.
54 plan: LirRelationExpr,
55 /// The arrangement keys that the plan is certain to produce.
56 keys: AvailableCollections,
57 /// Whether the plan's output may contain updates at future timestamps,
58 /// e.g., from a temporal MFP using `mz_now()`.
59 has_future_updates: bool,
60}
61
62pub(super) struct Context {
63 /// Known bindings to (possibly arranged) collections.
64 arrangements: BTreeMap<Id, AvailableCollections>,
65 /// Ids whose collections may contain updates at future timestamps,
66 /// e.g., from a temporal MFP using `mz_now()`.
67 has_future_updates: BTreeSet<Id>,
68 /// Tracks the next available `LirId`.
69 next_lir_id: LirId,
70 /// Information to print along with error messages.
71 debug_info: LirDebugInfo,
72 /// Whether to enable fusion of MFPs in reductions.
73 enable_reduce_mfp_fusion: bool,
74 /// Metrics recorded during lowering, if any are being collected.
75 metrics: Option<LoweringMetrics>,
76 /// Whether the current expression is subject to single-time (one-shot
77 /// `SELECT`) monotonic operator selection.
78 ///
79 /// Lowering locks in which arrangements a node makes available, and that set
80 /// changes with the chosen operator variant (e.g. a monotonic `TopK`/`Reduce`
81 /// arranges differently than its non-monotonic form). So the variant must be
82 /// picked here, during lowering, rather than by a later rewrite that would
83 /// leave the already-computed `AvailableCollections` describing the wrong shape.
84 ///
85 /// Initialized from the dataflow's `is_single_time()` and forced to `false`
86 /// while lowering the recursive bindings of a `LetRec`, whose values are not
87 /// restricted to a single time.
88 single_time: bool,
89 /// Global ids of the dataflow's source imports.
90 source_imports: BTreeSet<GlobalId>,
91 /// MIR `MfpPlan`s pushed onto `Get::Collection` reads of source imports,
92 /// keyed by the `Get`'s `LirId`.
93 ///
94 /// `refine_source_mfps` identifies common parts across sibling reads and pushes the
95 /// common part into the shared source MFP. That pass runs on MIR because
96 /// only MIR can utter the `mz_now()` predicates that temporal bounds fold
97 /// into. Retaining the MIR form here lets it find common parts without
98 /// round-tripping the LIR plans back through MIR.
99 source_get_mfps: BTreeMap<LirId, MfpPlan<MirScalarExpr>>,
100}
101
102impl Context {
103 pub fn new(
104 debug_name: String,
105 features: &OptimizerFeatures,
106 metrics: Option<&LoweringMetrics>,
107 ) -> Self {
108 Self {
109 arrangements: Default::default(),
110 has_future_updates: Default::default(),
111 next_lir_id: LirId(1),
112 debug_info: LirDebugInfo {
113 debug_name,
114 id: GlobalId::Transient(0),
115 },
116 enable_reduce_mfp_fusion: features.enable_reduce_mfp_fusion,
117 metrics: metrics.cloned(),
118 // Set from the dataflow in `lower` before any expression is lowered.
119 single_time: false,
120 source_imports: Default::default(),
121 source_get_mfps: Default::default(),
122 }
123 }
124
125 fn allocate_lir_id(&mut self) -> LirId {
126 let id = self.next_lir_id;
127 self.next_lir_id = LirId(
128 self.next_lir_id
129 .0
130 .checked_add(1)
131 .expect("No LirId overflow"),
132 );
133 id
134 }
135
136 pub fn lower(
137 mut self,
138 desc: DataflowDescription<OptimizedMirRelationExpr>,
139 ) -> Result<DataflowDescription<LirRelationExpr>, String> {
140 // Sources might provide arranged forms of their data, in the future.
141 // Indexes provide arranged forms of their data.
142 for IndexImport {
143 desc: index_desc,
144 typ,
145 ..
146 } in desc.index_imports.values()
147 {
148 let key = lses_from_mses(&index_desc.key);
149 // TODO[btv] - We should be told the permutation by
150 // `index_desc`, and it should have been generated
151 // at the same point the thinning logic was.
152 //
153 // We should for sure do that soon, but it requires
154 // a bit of a refactor, so for now we just
155 // _assume_ that they were both generated by `permutation_for_arrangement`,
156 // and recover it here.
157 let (permutation, thinning) = permutation_for_arrangement(&key, typ.arity());
158 let index_keys = self
159 .arrangements
160 .entry(Id::Global(index_desc.on_id))
161 .or_insert_with(AvailableCollections::default);
162 index_keys.arranged.push((key, permutation, thinning));
163 }
164 for id in desc.source_imports.keys() {
165 self.arrangements
166 .entry(Id::Global(*id))
167 .or_insert_with(AvailableCollections::new_raw);
168 self.source_imports.insert(*id);
169 }
170
171 // One-shot `SELECT` dataflows run at a single time, which lets us select
172 // monotonic operator variants during lowering (see the `TopK` and `Reduce`
173 // arms), so that `AvailableCollections` reflect the final operator variant.
174 self.single_time = desc.is_single_time();
175
176 // Build each object in order, registering the arrangements it forms.
177 let mut objects_to_build = Vec::with_capacity(desc.objects_to_build.len());
178 for build in desc.objects_to_build {
179 self.debug_info.id = build.id;
180 let LoweredExpr {
181 plan,
182 keys,
183 has_future_updates,
184 } = self.lower_mir_expr(&build.plan)?;
185
186 self.arrangements.insert(Id::Global(build.id), keys);
187 if has_future_updates {
188 self.has_future_updates.insert(Id::Global(build.id));
189 }
190 objects_to_build.push(BuildDesc { id: build.id, plan });
191 }
192
193 let mut dataflow = DataflowDescription {
194 source_imports: desc.source_imports,
195 index_imports: desc.index_imports,
196 objects_to_build,
197 index_exports: desc.index_exports,
198 sink_exports: desc.sink_exports,
199 as_of: desc.as_of,
200 until: desc.until,
201 initial_storage_as_of: desc.initial_storage_as_of,
202 refresh_schedule: desc.refresh_schedule,
203 debug_name: desc.debug_name,
204 time_dependence: desc.time_dependence,
205 };
206
207 // Refining: identify the common parts in the MFPs pushed onto a
208 // source's reads and hoist the shared prefix into the source itself.
209 self.refine_source_mfps(&mut dataflow);
210
211 Ok(dataflow)
212 }
213
214 /// Identifies common parts of the `MapFilterProject`s pushed onto sibling `Get::Collection`
215 /// reads of each imported source, hoisting the shared prefix into the
216 /// source's own MFP.
217 ///
218 /// The reads' MFPs are lowering artifacts (MIR sees only `Get(GlobalId)`),
219 /// so this belongs in lowering. We run on MIR because only MIR can
220 /// utter the `mz_now()` predicates that temporal bounds fold into, so it
221 /// consumes the MIR `MfpPlan`s stashed in [`Self::source_get_mfps`] rather
222 /// than round-tripping the lowered LIR plans back through MIR.
223 fn refine_source_mfps(&mut self, dataflow: &mut DataflowDescription<LirRelationExpr>) {
224 for (source_id, source_import) in dataflow.source_imports.iter_mut() {
225 let source = &mut source_import.desc;
226 let source_id = *source_id;
227 let mut identity_present = false;
228
229 // Collect the MIR `MfpPlan`s pushed onto this source's
230 // `Get::Collection` reads. Folding their temporal bounds back into
231 // `mz_now()` predicates (`into_map_filter_project`) lets
232 // `extract_common`'s column remapping apply uniformly across the
233 // whole MFP. Also note identity reads, which block pushdown.
234 let mut taken: Vec<(LirId, MapFilterProject<MirScalarExpr>)> = Vec::new();
235 for build_desc in dataflow.objects_to_build.iter() {
236 let mut todo = vec![&build_desc.plan];
237 while let Some(expression) = todo.pop() {
238 let node = &expression.node;
239 if let LirRelationNode::Get { id, plan, .. } = node {
240 if *id == Id::Global(source_id) {
241 match plan {
242 GetPlan::Collection(_) => {
243 let mir_plan = self
244 .source_get_mfps
245 .remove(&expression.lir_id)
246 .expect("stashed MIR MfpPlan for source Get::Collection");
247 taken.push((
248 expression.lir_id,
249 mir_plan.into_map_filter_project(),
250 ));
251 }
252 GetPlan::PassArrangements => {
253 identity_present = true;
254 }
255 GetPlan::Arrangement(..) => {
256 panic!("Surprising `GetPlan` for imported source: {:?}", plan);
257 }
258 }
259 }
260 } else {
261 todo.extend(node.children());
262 }
263 }
264 }
265
266 // Direct exports of sources are possible, and prevent pushdown.
267 identity_present |= dataflow
268 .index_exports
269 .values()
270 .any(|(x, _)| x.on_id == source_id);
271 identity_present |= dataflow.sink_exports.values().any(|x| x.from == source_id);
272
273 if identity_present || taken.is_empty() {
274 // Nothing to push down. The reads already carry their final LIR
275 // MFPs, so leave them untouched.
276 continue;
277 }
278
279 // Extract the common prefix and push it into the source's MFP.
280 let mut mfp_refs: Vec<&mut MapFilterProject<MirScalarExpr>> =
281 taken.iter_mut().map(|(_, mfp)| mfp).collect();
282 let common = MapFilterProject::extract_common(&mut mfp_refs[..]);
283 let mut source_mfp = if let Some(mfp) = source.arguments.operators.take() {
284 MapFilterProject::compose(mfp, common)
285 } else {
286 common
287 };
288 source_mfp.optimize();
289 source.arguments.operators = Some(source_mfp);
290
291 // Convert each residual MFP back to an LIR `MfpPlan` once, and
292 // install it on the corresponding read by `LirId`.
293 let replacements: BTreeMap<LirId, MfpPlan<LirScalarExpr>> = taken
294 .into_iter()
295 .map(|(lir_id, mir_mfp)| (lir_id, mfp_mir_to_lir_plan(mir_mfp)))
296 .collect();
297
298 for build_desc in dataflow.objects_to_build.iter_mut() {
299 let mut todo = vec![&mut build_desc.plan];
300 while let Some(expression) = todo.pop() {
301 if let Some(replacement) = replacements.get(&expression.lir_id) {
302 if let LirRelationNode::Get {
303 plan: GetPlan::Collection(mfp_plan),
304 ..
305 } = &mut expression.node
306 {
307 *mfp_plan = replacement.clone();
308 } else {
309 panic!(
310 "LirId {:?} was a GetPlan::Collection but is now {:?}",
311 expression.lir_id, expression.node
312 );
313 }
314 }
315 todo.extend(expression.node.children_mut());
316 }
317 }
318 }
319 }
320
321 /// This method converts a MirRelationExpr into a plan that can be directly rendered.
322 ///
323 /// The rough structure is that we repeatedly extract map/filter/project operators
324 /// from each expression we see, bundle them up as a `MapFilterProject` object, and
325 /// then produce a plan for the combination of that with the next operator.
326 ///
327 /// The method accesses `self.arrangements`, which it will locally add to and remove from for
328 /// `Let` bindings (by the end of the call it should contain the same bindings as when it
329 /// started).
330 ///
331 /// The result of the method is both a `LirRelationExpr`, but also a list of arrangements that
332 /// are certain to be produced, which can be relied on by the next steps in the plan.
333 /// Each of the arrangement keys is associated with an MFP that must be applied if that
334 /// arrangement is used, to back out the permutation associated with that arrangement.
335 ///
336 /// An empty list of arrangement keys indicates that only a `Collection` stream can
337 /// be assumed to exist.
338 fn lower_mir_expr(&mut self, expr: &MirRelationExpr) -> Result<LoweredExpr, String> {
339 // This function is recursive and can overflow its stack, so grow it if
340 // needed. The growth here is unbounded. Our general solution for this problem
341 // is to use [`ore::stack::RecursionGuard`] to additionally limit the stack
342 // depth. That however requires upstream error handling. This function is
343 // currently called by the Coordinator after calls to `catalog_transact`,
344 // and thus are not allowed to fail. Until that allows errors, we choose
345 // to allow the unbounded growth here. We are though somewhat protected by
346 // higher levels enforcing their own limits on stack depth (in the parser,
347 // transformer/desugarer, and planner).
348 mz_ore::stack::maybe_grow(|| self.lower_mir_expr_stack_safe(expr))
349 }
350
351 fn lower_mir_expr_stack_safe(&mut self, expr: &MirRelationExpr) -> Result<LoweredExpr, String> {
352 // Extract a maximally large MapFilterProject from `expr`.
353 // We will then try and push this in to the resulting expression.
354 //
355 // Importantly, `mfp` may contain temporal operators and not be a "safe" MFP.
356 // While we would eventually like all plan stages to be able to absorb such
357 // general operators, not all of them can.
358 let (mut mfp, expr) = MapFilterProject::extract_from_expression(expr);
359 // We attempt to plan what we have remaining, in the context of `mfp`.
360 // We may not be able to do this, and must wrap some operators with a `Mfp` stage.
361 let LoweredExpr {
362 mut plan,
363 mut keys,
364 mut has_future_updates,
365 } = match expr {
366 // These operators should have been extracted from the expression.
367 MirRelationExpr::Map { .. } => {
368 panic!("This operator should have been extracted");
369 }
370 MirRelationExpr::Filter { .. } => {
371 panic!("This operator should have been extracted");
372 }
373 MirRelationExpr::Project { .. } => {
374 panic!("This operator should have been extracted");
375 }
376 // These operators may not have been extracted, and need to result in a `LirRelationExpr`.
377 MirRelationExpr::Constant { rows, typ: _ } => {
378 let lir_id = self.allocate_lir_id();
379 let node = LirRelationNode::Constant {
380 rows: rows.clone().map(|rows| {
381 rows.into_iter()
382 .map(|(row, diff)| (StableRow(row), Timestamp::MIN, diff))
383 .collect()
384 }),
385 };
386 // The plan, not arranged in any way.
387 LoweredExpr {
388 plan: node.as_plan(lir_id),
389 keys: AvailableCollections::new_raw(),
390 has_future_updates: false,
391 }
392 }
393 MirRelationExpr::Get { id, typ: _, .. } => {
394 // This stage can absorb arbitrary MFP operators.
395 let mut mfp = mfp.take();
396 // If `mfp` is the identity, we can surface all imported arrangements.
397 // Otherwise, we apply `mfp` and promise no arrangements.
398 let mut in_keys = self
399 .arrangements
400 .get(id)
401 .cloned()
402 .unwrap_or_else(AvailableCollections::new_raw);
403
404 // Seek out an arrangement key that might be constrained to a literal.
405 // Note: this code has very little use nowadays, as its job was mostly taken over
406 // by `LiteralConstraints` (see in the below longer comment).
407 let key_val = in_keys
408 .arranged
409 .iter()
410 .filter_map(|key| {
411 mfp.literal_constraints(
412 &key.0.iter().map(MirScalarExpr::from).collect_vec(),
413 )
414 .map(|val| {
415 if let Some(metrics) = &self.metrics {
416 metrics.inc_literal_constraints("get");
417 }
418 (key.clone(), val)
419 })
420 })
421 .max_by_key(|(key, _val)| key.0.len());
422
423 // A source-import `Get::Collection`'s MIR `MfpPlan`, retained for
424 // `refine_source_mfps`. Stashed by `LirId` once the id is allocated below.
425 let mut source_get_mfp: Option<MfpPlan<MirScalarExpr>> = None;
426
427 // Determine the plan of action for the `Get` stage.
428 let plan = if let Some(((key, permutation, thinning), val)) = &key_val {
429 // This code path used to handle looking up literals from indexes, but it's
430 // mostly deprecated, as this is nowadays performed by the `LiteralConstraints`
431 // MIR transform instead. However, it's still called in a couple of tricky
432 // special cases:
433 // - `LiteralConstraints` handles only Gets of global ids, so this code still
434 // gets to handle Filters on top of Gets of local ids.
435 // - Lowering does a `MapFilterProject::extract_from_expression`, while
436 // `LiteralConstraints` does
437 // `MapFilterProject::extract_non_errors_from_expr_mut`.
438 // - It might happen that new literal constraint optimization opportunities
439 // appear somewhere near the end of the MIR optimizer after
440 // `LiteralConstraints` has already run.
441 // (Also note that a similar literal constraint handling machinery is also
442 // present when handling the leftover MFP after this big match.)
443 mfp.permute_fn(|c| permutation[c], thinning.len() + key.len());
444 in_keys.arranged = vec![(key.clone(), permutation.clone(), thinning.clone())];
445 GetPlan::Arrangement(
446 key.clone(),
447 Some(StableRow(val.clone())),
448 mfp_mir_to_lir_plan(mfp),
449 )
450 } else if !mfp.is_identity() {
451 // We need to ensure a collection exists, which means we must form it.
452 if let Some((key, permutation, thinning)) =
453 in_keys.arbitrary_arrangement().cloned()
454 {
455 mfp.permute_fn(|c| permutation[c], thinning.len() + key.len());
456 in_keys.arranged =
457 vec![(key.clone(), permutation.clone(), thinning.clone())];
458 GetPlan::Arrangement(key.clone(), None, mfp_mir_to_lir_plan(mfp))
459 } else {
460 let mir_plan = mfp.into_plan().expect("MFP planning failed");
461 if let Id::Global(gid) = id {
462 if self.source_imports.contains(gid) {
463 source_get_mfp = Some(mir_plan.clone());
464 }
465 }
466 GetPlan::Collection(mfp_plan_mir_to_lir(mir_plan))
467 }
468 } else {
469 // By default, just pass input arrangements through.
470 GetPlan::PassArrangements
471 };
472
473 let out_keys = if let GetPlan::PassArrangements = plan {
474 in_keys.clone()
475 } else {
476 AvailableCollections::new_raw()
477 };
478
479 // Even with a non-temporal MFP, we must propagate `has_future_updates`
480 // from the underlying binding — applying an MFP doesn't drop future-
481 // timestamped updates that already exist on the input.
482 //
483 // Note that global Gets from different dataflows can't have future updates, because
484 // both indexes and materialized views hold back future updates.
485 let has_future_updates = self.has_future_updates.contains(id)
486 || match &plan {
487 GetPlan::Arrangement(_, _, mfp_plan) | GetPlan::Collection(mfp_plan) => {
488 mfp_plan.has_temporal_bounds()
489 }
490 GetPlan::PassArrangements => false,
491 };
492
493 let lir_id = self.allocate_lir_id();
494 if let Some(mir_plan) = source_get_mfp {
495 self.source_get_mfps.insert(lir_id, mir_plan);
496 }
497 let node = LirRelationNode::Get {
498 id: id.clone(),
499 keys: in_keys,
500 plan,
501 };
502 // Return the plan, and any keys if an identity `mfp`.
503 LoweredExpr {
504 plan: node.as_plan(lir_id),
505 keys: out_keys,
506 has_future_updates,
507 }
508 }
509 MirRelationExpr::Let { id, value, body } => {
510 // It would be unfortunate to have a non-trivial `mfp` here, as we hope
511 // that they would be pushed down. I am not sure if we should take the
512 // initiative to push down the `mfp` ourselves.
513
514 // Plan the value using only the initial arrangements, but
515 // introduce any resulting arrangements bound to `id`.
516 let LoweredExpr {
517 plan: value,
518 keys: v_keys,
519 has_future_updates: v_future,
520 } = self.lower_mir_expr(value)?;
521 let pre_existing = self.arrangements.insert(Id::Local(*id), v_keys);
522 assert_none!(pre_existing);
523 if v_future {
524 self.has_future_updates.insert(Id::Local(*id));
525 }
526 // Plan the body using initial and `value` arrangements,
527 // and then remove reference to the value arrangements.
528 let LoweredExpr {
529 plan: body,
530 keys: b_keys,
531 has_future_updates: b_future,
532 } = self.lower_mir_expr(body)?;
533 self.arrangements.remove(&Id::Local(*id));
534 self.has_future_updates.remove(&Id::Local(*id));
535 // Return the plan, and any `body` arrangements.
536 let lir_id = self.allocate_lir_id();
537 LoweredExpr {
538 plan: LirRelationNode::Let {
539 id: id.clone(),
540 value: Box::new(value),
541 body: Box::new(body),
542 }
543 .as_plan(lir_id),
544 keys: b_keys,
545 has_future_updates: b_future,
546 }
547 }
548 MirRelationExpr::LetRec {
549 ids,
550 values,
551 limits,
552 body,
553 } => {
554 assert_eq!(ids.len(), values.len());
555 assert_eq!(ids.len(), limits.len());
556 // Plan the values using only the available arrangements, but
557 // introduce any resulting arrangements bound to each `id`.
558 // Arrangements made available cannot be used by prior bindings,
559 // as we cannot circulate an arrangement through a `Variable` yet.
560 let mut lir_values = Vec::with_capacity(values.len());
561 let mut any_v_future = false;
562 // The recursive bindings of a `LetRec` are not restricted to a single
563 // time, so single-time monotonic selection must not apply to them. Only
564 // the `body`, lowered below, inherits the enclosing scope's flag.
565 let outer_single_time = self.single_time;
566 self.single_time = false;
567 for (id, value) in ids.iter().zip_eq(values) {
568 let LoweredExpr {
569 plan: mut lir_value,
570 keys: mut v_keys,
571 has_future_updates: v_future,
572 } = self.lower_mir_expr(value)?;
573 any_v_future |= v_future;
574 // If `v_keys` does not contain an unarranged collection, we must form it.
575 if !v_keys.raw {
576 // Choose an "arbitrary" arrangement; TODO: prefer a specific one.
577 let (input_key, permutation, thinning) =
578 v_keys.arbitrary_arrangement().unwrap();
579 let mut input_mfp = MapFilterProject::new(value.arity());
580 input_mfp.permute_fn(|c| permutation[c], thinning.len() + input_key.len());
581 let input_key = Some(input_key.clone());
582
583 let forms = AvailableCollections::new_raw();
584
585 // We just want to insert an `ArrangeBy` to form an unarranged collection,
586 // but there is a complication: We shouldn't break the invariant (created by
587 // `NormalizeLets`, and relied upon by the rendering) that there isn't
588 // anything between two `LetRec`s. So if `lir_value` is itself a `LetRec`,
589 // then we insert the `ArrangeBy` on the `body` of the inner `LetRec`,
590 // instead of on top of the inner `LetRec`.
591 //
592 // We forward `v_future` for honesty; bucketing has no observable effect
593 // inside an iterative scope, but the field should reflect reality.
594 lir_value = match lir_value {
595 LirRelationExpr {
596 node:
597 LirRelationNode::LetRec {
598 ids,
599 values,
600 limits,
601 body,
602 },
603 lir_id,
604 } => {
605 let inner_lir_id = self.allocate_lir_id();
606 LirRelationNode::LetRec {
607 ids,
608 values,
609 limits,
610 body: Box::new(
611 LirRelationNode::ArrangeBy {
612 input_key,
613 input: body,
614 input_mfp: mfp_mir_to_lir_plan(input_mfp),
615 forms,
616 strategy: strategy_from_future(v_future),
617 }
618 .as_plan(inner_lir_id),
619 ),
620 }
621 .as_plan(lir_id)
622 }
623 lir_value => {
624 let lir_id = self.allocate_lir_id();
625 LirRelationNode::ArrangeBy {
626 input_key,
627 input: Box::new(lir_value),
628 input_mfp: mfp_mir_to_lir_plan(input_mfp),
629 forms,
630 strategy: strategy_from_future(v_future),
631 }
632 .as_plan(lir_id)
633 }
634 };
635 v_keys.raw = true;
636 }
637 let pre_existing = self.arrangements.insert(Id::Local(*id), v_keys);
638 assert_none!(pre_existing);
639 if v_future {
640 self.has_future_updates.insert(Id::Local(*id));
641 }
642 lir_values.push(lir_value);
643 }
644 // As we exit the iterative scope, we must leave all arrangements behind,
645 // as they reference a timestamp coordinate that must be stripped off.
646 for id in ids.iter() {
647 self.arrangements
648 .insert(Id::Local(*id), AvailableCollections::new_raw());
649 }
650 // Plan the body using initial and `value` arrangements,
651 // and then remove reference to the value arrangements.
652 self.single_time = outer_single_time;
653 let LoweredExpr {
654 plan: body,
655 keys: b_keys,
656 has_future_updates: b_future,
657 } = self.lower_mir_expr(body)?;
658 for id in ids.iter() {
659 self.arrangements.remove(&Id::Local(*id));
660 self.has_future_updates.remove(&Id::Local(*id));
661 }
662 // Return the plan, and any `body` arrangements.
663 //
664 // The body's `b_future` alone can under-report: an earlier binding may only
665 // inherit `has_future_updates` via a Variable to a *later* binding, which the
666 // sequential sweep can't observe at the time the earlier binding is lowered.
667 // A precise fix would require a fixpoint (or the MIR `Analysis` framework with
668 // a `true ⊑ false` lattice). As a cheap correct alternative, OR with the
669 // bindings' future flags: any cross-binding propagation must originate from a
670 // local temporal predicate inside *some* binding, so the OR captures it
671 // without forcing bucketing on a fully non-temporal LetRec.
672 let lir_id = self.allocate_lir_id();
673 LoweredExpr {
674 plan: LirRelationNode::LetRec {
675 ids: ids.clone(),
676 values: lir_values,
677 limits: limits.clone(),
678 body: Box::new(body),
679 }
680 .as_plan(lir_id),
681 keys: b_keys,
682 has_future_updates: b_future || any_v_future,
683 }
684 }
685 MirRelationExpr::FlatMap {
686 input: flat_map_input,
687 func,
688 exprs,
689 } => {
690 // A `FlatMap UnnestList` that comes after the `Reduce` of a window function can be
691 // fused into the lowered `Reduce`.
692 //
693 // In theory, we could have implemented this also as an MIR transform. However, this
694 // is more of a physical optimization, which are sometimes unpleasant to make a part
695 // of the MIR pipeline. The specific problem here with putting this into the MIR
696 // pipeline would be that we'd need to modify MIR's semantics: MIR's Reduce
697 // currently always emits exactly 1 row per group, but the fused Reduce-FlatMap can
698 // emit multiple rows per group. Such semantic changes of MIR are very scary, since
699 // various parts of the optimizer assume that Reduce emits only 1 row per group, and
700 // it would be very hard to hunt down all these parts. (For example, key inference
701 // infers the group key as a unique key.)
702 let fused_with_reduce = 'fusion: {
703 if !matches!(func, TableFunc::UnnestList { .. }) {
704 break 'fusion None;
705 }
706 // We might have a Project of a single col between the FlatMap and the
707 // Reduce. (It projects away the grouping keys of the Reduce, and keeps the
708 // result of the window function.)
709 let (maybe_reduce, num_grouping_keys) = if let MirRelationExpr::Project {
710 input: project_input,
711 outputs: projection,
712 } = &**flat_map_input
713 {
714 // We want this to be a single column, because we'll want to deal with only
715 // one aggregation in the `Reduce`. (The aggregation of a window function
716 // always stands alone currently: we plan them separately from other
717 // aggregations, and Reduces are never fused. When window functions are
718 // fused with each other, they end up in one aggregation. When there are
719 // multiple window functions in the same SELECT, but can't be fused, they
720 // end up in different Reduces.)
721 if let &[single_col] = &**projection {
722 (project_input, single_col)
723 } else {
724 break 'fusion None;
725 }
726 } else {
727 (flat_map_input, 0)
728 };
729 if let MirRelationExpr::Reduce {
730 input,
731 group_key,
732 aggregates,
733 monotonic,
734 expected_group_size,
735 } = &**maybe_reduce
736 {
737 if group_key.len() != num_grouping_keys
738 || aggregates.len() != 1
739 || !aggregates[0].func.can_fuse_with_unnest_list()
740 {
741 break 'fusion None;
742 }
743 // At the beginning, `non_fused_mfp_above_flat_map` will be the original MFP
744 // above the FlatMap. Later, we'll mutate this to be the residual MFP that
745 // didn't get fused into the `Reduce`.
746 let non_fused_mfp_above_flat_map = &mut mfp;
747 let reduce_output_arity = num_grouping_keys + 1;
748 // We are fusing away the list that the FlatMap would have been unnesting,
749 // so the column that had that list disappears, so we have to permute the
750 // MFP above the FlatMap with this column disappearance.
751 let tweaked_mfp = {
752 let mut mfp = non_fused_mfp_above_flat_map.clone();
753 if mfp.demand().contains(&0) {
754 // I don't think this can happen currently that this MFP would
755 // refer to the list column, because both the list column and the
756 // MFP were constructed by the HIR-to-MIR lowering, so it's not just
757 // some random MFP that we are seeing here. But anyhow, it's better
758 // to check this here for robustness against future code changes.
759 break 'fusion None;
760 }
761 let permutation: BTreeMap<_, _> =
762 (1..mfp.input_arity).map(|col| (col, col - 1)).collect();
763 mfp.permute_fn(|c| permutation[&c], mfp.input_arity - 1);
764 mfp
765 };
766 // We now put together the project that was before the FlatMap, and the
767 // tweaked version of the MFP that was after the FlatMap.
768 // (Part of this MFP might be fused into the Reduce.)
769 let mut project_and_tweaked_mfp = {
770 let mut mfp = MapFilterProject::new(reduce_output_arity);
771 mfp = mfp.project(vec![num_grouping_keys]);
772 mfp = MapFilterProject::compose(mfp, tweaked_mfp);
773 mfp
774 };
775 let fused = self.lower_reduce(
776 input,
777 group_key,
778 aggregates,
779 monotonic,
780 expected_group_size,
781 &mut project_and_tweaked_mfp,
782 true,
783 )?;
784 // Update the residual MFP.
785 *non_fused_mfp_above_flat_map = project_and_tweaked_mfp;
786 Some(fused)
787 } else {
788 break 'fusion None;
789 }
790 };
791 if let Some(fused_with_reduce) = fused_with_reduce {
792 fused_with_reduce
793 } else {
794 // Couldn't fuse it with a `Reduce`, so lower as a normal `FlatMap`.
795 let LoweredExpr {
796 plan: input,
797 keys,
798 has_future_updates: input_future,
799 } = self.lower_mir_expr(flat_map_input)?;
800 // This stage can absorb arbitrary MFP instances.
801 let mut mfp = mfp.take();
802 let mut exprs = exprs.clone();
803 // Prefer the unarranged collection when present: it presents input columns
804 // in logical order, so no permutation is required.
805 let input_key = if keys.raw {
806 None
807 } else if let Some((k, permutation, thinning)) = keys.arbitrary_arrangement() {
808 // Reading from this arrangement exposes input columns in arrangement
809 // order (key columns followed by thinned value columns). We must
810 // permute every reference to an input column accordingly: the
811 // `expr`s feeding the table function arguments, and the `mfp` running
812 // after the table function (which still references input columns at
813 // positions `0..input_arity`).
814 //
815 // The renderer hands the `mfp` the *whole* arranged row and appends the
816 // table-function output after it. The arranged row can be wider than the
817 // logical input row when the key is not a set of distinct columns (an
818 // expression, functional, or repeated-column key carries extra key
819 // values). So the table-function output columns at positions
820 // `input_arity..` must be shifted to land after the arranged row, and the
821 // `mfp`'s new input arity must reflect the arranged width.
822 for expr in &mut exprs {
823 expr.permute(permutation);
824 }
825 let input_arity = permutation.len();
826 let arranged_arity = thinning.len() + k.len();
827 let output_arity = mfp.input_arity - input_arity;
828 mfp.permute_fn(
829 |c| {
830 if c < input_arity {
831 permutation[c]
832 } else {
833 arranged_arity + (c - input_arity)
834 }
835 },
836 arranged_arity + output_arity,
837 );
838 Some(k.clone())
839 } else {
840 None
841 };
842
843 let lir_id = self.allocate_lir_id();
844 // The absorbed `mfp` may contain temporal predicates, which can
845 // introduce future-stamped updates that aren't present on the input.
846 let has_future_updates = input_future || mfp.has_temporal_predicates();
847 // Return the plan, and no arrangements.
848 LoweredExpr {
849 plan: LirRelationNode::FlatMap {
850 input_key,
851 input: Box::new(input),
852 exprs: lses_from_mses(&exprs),
853 func: func.clone(),
854 mfp_after: mfp_mir_to_lir_plan(mfp),
855 }
856 .as_plan(lir_id),
857 keys: AvailableCollections::new_raw(),
858 has_future_updates,
859 }
860 }
861 }
862 MirRelationExpr::Join {
863 inputs,
864 equivalences,
865 implementation,
866 } => {
867 // Plan each of the join inputs independently.
868 // The `plans` get surfaced upwards, and the `input_keys` should
869 // be used as part of join planning / to validate the existing
870 // plans / to aid in indexed seeding of update streams.
871 let mut plans = Vec::new();
872 let mut input_keys = Vec::new();
873 let mut input_arities = Vec::new();
874 let mut input_futures = Vec::new();
875 for input in inputs.iter() {
876 let LoweredExpr {
877 plan,
878 keys,
879 has_future_updates: input_future,
880 } = self.lower_mir_expr(input)?;
881 input_arities.push(input.arity());
882 plans.push(plan);
883 input_keys.push(keys);
884 input_futures.push(input_future);
885 }
886 let any_input_future = input_futures.iter().any(|&f| f);
887
888 let input_mapper =
889 JoinInputMapper::new_from_input_arities(input_arities.iter().copied());
890
891 // Extract temporal predicates as joins cannot currently absorb them.
892 let (plan, missing) = match implementation {
893 IndexedFilter(_coll_id, _idx_id, key, _val) => {
894 // Start with the constant input. (This used to be important before database-issues#4016
895 // was fixed.)
896 let start: usize = 1;
897 let order = vec![(0usize, key.clone(), None)];
898 // All columns of the constant input will be part of the arrangement key.
899 let source_arrangement = (
900 (0..key.len())
901 .map(LirScalarExpr::column)
902 .collect::<Vec<_>>(),
903 (0..key.len()).collect::<Vec<_>>(),
904 Vec::<usize>::new(),
905 );
906 let (ljp, missing) = LinearJoinPlan::create_from(
907 start,
908 Some(&source_arrangement),
909 equivalences,
910 &order,
911 input_mapper,
912 &mut mfp,
913 &input_keys,
914 );
915 (JoinPlan::Linear(ljp), missing)
916 }
917 Differential((start, start_arr, _start_characteristic), order) => {
918 let source_arrangement = start_arr.as_ref().and_then(|key| {
919 let key = lses_from_mses(key);
920 input_keys[*start]
921 .arranged
922 .iter()
923 .find(|(k, _, _)| k == &key)
924 .clone()
925 });
926 let (ljp, missing) = LinearJoinPlan::create_from(
927 *start,
928 source_arrangement,
929 equivalences,
930 order,
931 input_mapper,
932 &mut mfp,
933 &input_keys,
934 );
935 (JoinPlan::Linear(ljp), missing)
936 }
937 DeltaQuery(orders) => {
938 let (djp, missing) = DeltaJoinPlan::create_from(
939 equivalences,
940 orders,
941 input_mapper,
942 &mut mfp,
943 &input_keys,
944 );
945 (JoinPlan::Delta(djp), missing)
946 }
947 // Other plans are errors, and should be reported as such.
948 Unimplemented => return Err("unimplemented join".to_string()),
949 };
950 // The renderer will expect certain arrangements to exist; if any of those are not available, the join planning functions above should have returned them in
951 // `missing`. We thus need to plan them here so they'll exist.
952 let is_delta = matches!(plan, JoinPlan::Delta(_));
953 for ((((input_plan, input_keys), missing), arity), input_future) in plans
954 .iter_mut()
955 .zip_eq(input_keys.iter())
956 .zip_eq(missing)
957 .zip_eq(input_arities.iter().cloned())
958 .zip_eq(input_futures.iter().copied())
959 {
960 if missing != Default::default() {
961 if is_delta {
962 // join_implementation.rs produced a sub-optimal plan here;
963 // we shouldn't plan delta joins at all if not all of the required
964 // arrangements are available. Soft panic in CI and log an error in
965 // production to increase the chances that we will catch all situations
966 // that violate this constraint.
967 soft_panic_or_log!("Arrangements depended on by delta join alarmingly absent: {:?}
968Dataflow info: {}
969This is not expected to cause incorrect results, but could indicate a performance issue in Materialize.", missing, self.debug_info);
970 } else {
971 soft_panic_or_log!("Arrangements depended on by a non-delta join are absent: {:?}
972Dataflow info: {}
973This is not expected to cause incorrect results, but could indicate a performance issue in Materialize.", missing, self.debug_info);
974 // Nowadays MIR transforms take care to insert MIR ArrangeBys for each
975 // Join input. (Earlier, they were missing in the following cases:
976 // - They were const-folded away for constant inputs. This is not
977 // happening since
978 // https://github.com/MaterializeInc/materialize/pull/16351
979 // - They were not being inserted for the constant input of
980 // `IndexedFilter`s. This was fixed in
981 // https://github.com/MaterializeInc/materialize/pull/20920
982 // - They were not being inserted for the first input of Differential
983 // joins. This was fixed in
984 // https://github.com/MaterializeInc/materialize/pull/16099)
985 }
986 let lir_id = self.allocate_lir_id();
987 let raw_plan = std::mem::replace(
988 input_plan,
989 LirRelationNode::Constant {
990 rows: Ok(Vec::new()),
991 }
992 .as_plan(lir_id),
993 );
994 *input_plan =
995 self.arrange_by(raw_plan, missing, input_keys, arity, input_future);
996 }
997 }
998 // Return the plan, and no arrangements.
999 // Both linear and delta join planning extract temporal predicates back into the
1000 // residual `mfp` (see `LinearJoinPlan::create_from` / `DeltaJoinPlan::create_from`),
1001 // so the absorbed MFP cannot introduce future updates — the join's output future
1002 // flag is just the OR of its inputs.
1003 let lir_id = self.allocate_lir_id();
1004 LoweredExpr {
1005 plan: LirRelationNode::Join {
1006 inputs: plans,
1007 plan,
1008 }
1009 .as_plan(lir_id),
1010 keys: AvailableCollections::new_raw(),
1011 has_future_updates: any_input_future,
1012 }
1013 }
1014 MirRelationExpr::Reduce {
1015 input,
1016 group_key,
1017 aggregates,
1018 monotonic,
1019 expected_group_size,
1020 } => {
1021 if aggregates
1022 .iter()
1023 .any(|agg| agg.func.can_fuse_with_unnest_list())
1024 {
1025 // This case should have been handled at the `MirRelationExpr::FlatMap` case
1026 // above. But that has a pretty complicated pattern matching, so it's not
1027 // unthinkable that it fails.
1028 soft_panic_or_log!(
1029 "Window function performance issue: `reduce_unnest_list_fusion` failed"
1030 );
1031 }
1032 self.lower_reduce(
1033 input,
1034 group_key,
1035 aggregates,
1036 monotonic,
1037 expected_group_size,
1038 &mut mfp,
1039 false,
1040 )?
1041 }
1042 MirRelationExpr::TopK {
1043 input,
1044 group_key,
1045 order_key,
1046 limit,
1047 offset,
1048 monotonic,
1049 expected_group_size,
1050 } => {
1051 let arity = input.arity();
1052 let LoweredExpr {
1053 plan: input,
1054 keys,
1055 has_future_updates: input_future,
1056 } = self.lower_mir_expr(input)?;
1057
1058 let mut top_k_plan = TopKPlan::create_from(
1059 group_key.clone(),
1060 order_key.clone(),
1061 *offset,
1062 limit
1063 .as_ref()
1064 .map(|limit| LirScalarExpr::try_from(limit).expect("lowerable MIR")),
1065 arity,
1066 *monotonic,
1067 *expected_group_size,
1068 );
1069
1070 // For single-time dataflows, upgrade to the monotonic variant with
1071 // mandatory consolidation. `refine_single_time_consolidation` later
1072 // relaxes `must_consolidate` where the input is physically monotonic.
1073 if self.single_time {
1074 top_k_plan.as_monotonic(true);
1075 }
1076
1077 // We don't have an MFP here -- install an operator to permute the
1078 // input, if necessary.
1079 let input = if !keys.raw {
1080 self.arrange_by(
1081 input,
1082 AvailableCollections::new_raw(),
1083 &keys,
1084 arity,
1085 // `new_raw` means no arrangement, so no bucketing is needed
1086 false,
1087 )
1088 } else {
1089 input
1090 };
1091 // Return the plan, and the keys it produces. `MonotonicTop1` arranges its
1092 // output by the group key (see `render_top1_monotonic`), so a downstream
1093 // consumer keyed the same way can reuse that arrangement instead of forcing
1094 // another `ArrangeBy`.
1095 let out_keys = match &top_k_plan {
1096 TopKPlan::MonotonicTop1(_) => {
1097 let key = group_key
1098 .iter()
1099 .map(|c| LirScalarExpr::column(*c))
1100 .collect::<Vec<_>>();
1101 let (permutation, thinning) = permutation_for_arrangement(&key, arity);
1102 AvailableCollections::new_arranged(vec![(key, permutation, thinning)])
1103 }
1104 // MonotonicTopK / Basic key their arrangements by (hash, group_key), which is
1105 // not reusable by a group-key consumer, so they advertise no arrangement.
1106 TopKPlan::MonotonicTopK(_) | TopKPlan::Basic(_) => {
1107 AvailableCollections::new_raw()
1108 }
1109 };
1110 let temporal_bucketing_strategy = strategy_from_future(input_future);
1111 let lir_id = self.allocate_lir_id();
1112 LoweredExpr {
1113 plan: LirRelationNode::TopK {
1114 input: Box::new(input),
1115 top_k_plan,
1116 temporal_bucketing_strategy,
1117 }
1118 .as_plan(lir_id),
1119 keys: out_keys,
1120 has_future_updates: false,
1121 }
1122 }
1123 MirRelationExpr::Negate { input } => {
1124 let arity = input.arity();
1125 let LoweredExpr {
1126 plan: input,
1127 keys,
1128 has_future_updates: input_future,
1129 } = self.lower_mir_expr(input)?;
1130
1131 // We don't have an MFP here -- install an operator to permute the
1132 // input, if necessary.
1133 let input = if !keys.raw {
1134 self.arrange_by(
1135 input,
1136 AvailableCollections::new_raw(),
1137 &keys,
1138 arity,
1139 // `new_raw` means no arrangement, so no bucketing is needed
1140 false,
1141 )
1142 } else {
1143 input
1144 };
1145 // Return the plan, and no arrangements.
1146 let lir_id = self.allocate_lir_id();
1147 LoweredExpr {
1148 plan: LirRelationNode::Negate {
1149 input: Box::new(input),
1150 }
1151 .as_plan(lir_id),
1152 keys: AvailableCollections::new_raw(),
1153 has_future_updates: input_future,
1154 }
1155 }
1156 MirRelationExpr::Threshold { input } => {
1157 let LoweredExpr {
1158 plan,
1159 keys,
1160 has_future_updates: input_future,
1161 } = self.lower_mir_expr(input)?;
1162 let arity = input.arity();
1163 let (threshold_plan, required_arrangement) = ThresholdPlan::create_from(arity);
1164
1165 let plan = if !keys
1166 .arranged
1167 .iter()
1168 .any(|(key, _, _)| key == &required_arrangement.0)
1169 {
1170 self.arrange_by(
1171 plan,
1172 AvailableCollections::new_arranged(vec![required_arrangement]),
1173 &keys,
1174 arity,
1175 input_future,
1176 )
1177 } else {
1178 plan
1179 };
1180
1181 let output_keys = threshold_plan.keys();
1182 // Return the plan, and any produced keys.
1183 let lir_id = self.allocate_lir_id();
1184 LoweredExpr {
1185 plan: LirRelationNode::Threshold {
1186 input: Box::new(plan),
1187 threshold_plan,
1188 }
1189 .as_plan(lir_id),
1190 keys: output_keys,
1191 // Threshold builds its own output arrangement whose
1192 // MergeBatcher absorbs future-stamped updates, so no
1193 // future updates flow out.
1194 has_future_updates: false,
1195 }
1196 }
1197 MirRelationExpr::Union { base, inputs } => {
1198 let arity = base.arity();
1199 let mut lowered_inputs = Vec::with_capacity(1 + inputs.len());
1200 lowered_inputs.push(self.lower_mir_expr(base)?);
1201 for input in inputs.iter() {
1202 lowered_inputs.push(self.lower_mir_expr(input)?);
1203 }
1204
1205 // A Union with any `Negate` input should consolidate its
1206 // output. The lowering is the only place where this decision
1207 // can be coupled with the per-input bucketing strategy.
1208 let consolidate_output = lowered_inputs
1209 .iter()
1210 .any(|l| matches!(l.plan.node, LirRelationNode::Negate { .. }));
1211
1212 // Per-input bucketing strategies: only meaningful when the
1213 // Union consolidates its output, since bucketing only pays off
1214 // ahead of a downstream consolidator.
1215 let temporal_bucketing_strategies: Vec<ArrangementStrategy> = if consolidate_output
1216 {
1217 lowered_inputs
1218 .iter()
1219 .map(|l| strategy_from_future(l.has_future_updates))
1220 .collect()
1221 } else {
1222 lowered_inputs
1223 .iter()
1224 .map(|_| ArrangementStrategy::Direct)
1225 .collect()
1226 };
1227
1228 let has_future_updates = if consolidate_output {
1229 // The MergeBatcher will hold back future updates (regardless of whether we are
1230 // bucketing here or not).
1231 false
1232 } else {
1233 lowered_inputs.iter().any(|l| l.has_future_updates)
1234 };
1235
1236 let plans = lowered_inputs
1237 .into_iter()
1238 .map(
1239 |LoweredExpr {
1240 plan,
1241 keys,
1242 has_future_updates: _,
1243 }| {
1244 // We don't have an MFP here -- install an operator to permute the
1245 // input, if necessary.
1246 if !keys.raw {
1247 self.arrange_by(
1248 plan,
1249 AvailableCollections::new_raw(),
1250 &keys,
1251 arity,
1252 // `new_raw` means no arrangement, so no bucketing is needed
1253 false,
1254 )
1255 } else {
1256 plan
1257 }
1258 },
1259 )
1260 .collect();
1261 // Return the plan and no arrangements.
1262 let lir_id = self.allocate_lir_id();
1263 LoweredExpr {
1264 plan: LirRelationNode::Union {
1265 inputs: plans,
1266 consolidate_output,
1267 temporal_bucketing_strategies,
1268 }
1269 .as_plan(lir_id),
1270 keys: AvailableCollections::new_raw(),
1271 has_future_updates,
1272 }
1273 }
1274 MirRelationExpr::ArrangeBy { input, keys } => {
1275 let input_mir = input;
1276 let LoweredExpr {
1277 plan: input,
1278 keys: mut input_keys,
1279 has_future_updates: input_has_future_updates,
1280 } = self.lower_mir_expr(input)?;
1281 // Fill the `types` in `input_keys` if not already present.
1282 let arity = input_mir.arity();
1283
1284 // Determine keys that are not present in `input_keys`.
1285 let new_keys = keys
1286 .iter()
1287 .filter(|k1| {
1288 !input_keys.arranged.iter().any(|(k2, _, _)| {
1289 k1.len() == k2.len()
1290 && k1
1291 .iter()
1292 .zip_eq(k2)
1293 .all(|(e1, e2)| *e1 == MirScalarExpr::from(e2))
1294 })
1295 })
1296 .cloned()
1297 .collect::<Vec<_>>();
1298 if new_keys.is_empty() {
1299 LoweredExpr {
1300 plan: input,
1301 keys: input_keys,
1302 has_future_updates: input_has_future_updates,
1303 }
1304 } else {
1305 let mut new_keys = new_keys
1306 .iter()
1307 .map(|k| {
1308 let k = lses_from_mses(k);
1309 let (permutation, thinning) = permutation_for_arrangement(&k, arity);
1310 (k, permutation, thinning)
1311 })
1312 .collect::<Vec<_>>();
1313 let forms = AvailableCollections {
1314 raw: input_keys.raw,
1315 arranged: new_keys.clone(),
1316 };
1317 let (input_key, input_mfp) = if let Some((input_key, permutation, thinning)) =
1318 input_keys.arbitrary_arrangement()
1319 {
1320 let mut mfp = MapFilterProject::new(arity);
1321 mfp.permute_fn(|c| permutation[c], thinning.len() + input_key.len());
1322 (Some(input_key.clone()), mfp)
1323 } else {
1324 (None, MapFilterProject::new(arity))
1325 };
1326 input_keys.arranged.append(&mut new_keys);
1327 input_keys.arranged.sort_by(|k1, k2| k1.0.cmp(&k2.0));
1328
1329 // Return the plan and extended keys.
1330 let lir_id = self.allocate_lir_id();
1331 let strategy = strategy_from_future(input_has_future_updates);
1332 assert!(!forms.arranged.is_empty()); // i.e., we do build an arrangement
1333 let has_future_updates = false;
1334 LoweredExpr {
1335 plan: LirRelationNode::ArrangeBy {
1336 input_key,
1337 input: Box::new(input),
1338 input_mfp: mfp_mir_to_lir_plan(input_mfp),
1339 forms,
1340 strategy,
1341 }
1342 .as_plan(lir_id),
1343 keys: input_keys,
1344 has_future_updates,
1345 }
1346 }
1347 }
1348 };
1349
1350 // If the plan stage did not absorb all linear operators, introduce a new stage to implement them.
1351 if !mfp.is_identity() {
1352 // Check if this MFP introduces future updates.
1353 let mfp_is_temporal = mfp.has_temporal_predicates();
1354 has_future_updates = has_future_updates || mfp_is_temporal;
1355 // Seek out an arrangement key that might be constrained to a literal.
1356 // TODO: Improve key selection heuristic.
1357 let key_val = keys
1358 .arranged
1359 .iter()
1360 .filter_map(|(key, permutation, thinning)| {
1361 let mut mfp = mfp.clone();
1362 mfp.permute_fn(|c| permutation[c], thinning.len() + key.len());
1363 mfp.literal_constraints(&key.iter().map(MirScalarExpr::from).collect_vec())
1364 .map(|val| {
1365 if let Some(metrics) = &self.metrics {
1366 metrics.inc_literal_constraints("mfp");
1367 }
1368 (key.clone(), permutation, thinning, val)
1369 })
1370 })
1371 .max_by_key(|(key, _, _, _)| key.len());
1372
1373 // Input key selection strategy:
1374 // (1) If we can read a key at a particular value, do so
1375 // (2) Otherwise, if there is a key that causes the MFP to be the identity, and
1376 // therefore allows us to avoid discarding the arrangement, use that.
1377 // (3) Otherwise, if there is _some_ key, use that,
1378 // (4) Otherwise just read the raw collection.
1379 let input_key_val = if let Some((key, permutation, thinning, val)) = key_val {
1380 mfp.permute_fn(|c| permutation[c], thinning.len() + key.len());
1381
1382 Some((key, Some(val)))
1383 } else if let Some((key, permutation, thinning)) =
1384 keys.arranged.iter().find(|(key, permutation, thinning)| {
1385 let mut mfp = mfp.clone();
1386 mfp.permute_fn(|c| permutation[c], thinning.len() + key.len());
1387 mfp.is_identity()
1388 })
1389 {
1390 mfp.permute_fn(|c| permutation[c], thinning.len() + key.len());
1391 Some((key.clone(), None))
1392 } else if let Some((key, permutation, thinning)) = keys.arbitrary_arrangement() {
1393 mfp.permute_fn(|c| permutation[c], thinning.len() + key.len());
1394 Some((key.clone(), None))
1395 } else {
1396 None
1397 };
1398
1399 if mfp.is_identity() {
1400 // We have discovered a key
1401 // whose permutation causes the MFP to actually
1402 // be the identity! We can keep it around,
1403 // but without its permutation this time,
1404 // and with a trivial thinning of the right length.
1405 let (key, val) = input_key_val.unwrap();
1406 let (_old_key, old_permutation, old_thinning) = keys
1407 .arranged
1408 .iter_mut()
1409 .find(|(key2, _, _)| key2 == &key)
1410 .unwrap();
1411 *old_permutation = (0..mfp.input_arity).collect();
1412 let old_thinned_arity = old_thinning.len();
1413 *old_thinning = (0..old_thinned_arity).collect();
1414 // Get rid of all other forms, as this is now the only one known to be valid.
1415 // TODO[btv] we can probably save the other arrangements too, if we adjust their permutations.
1416 // This is not hard to do, but leaving it for a quick follow-up to avoid making the present diff too unwieldy.
1417 keys.arranged.retain(|(key2, _, _)| key2 == &key);
1418 keys.raw = false;
1419
1420 // Creating a LirRelationExpr::Mfp node is now logically unnecessary, but we
1421 // should do so anyway when `val` is populated, so that
1422 // the `key_val` optimization gets applied.
1423 let lir_id = self.allocate_lir_id();
1424 if val.is_some() {
1425 plan = LirRelationNode::Mfp {
1426 input: Box::new(plan),
1427 mfp: mfp_mir_to_lir_plan(mfp),
1428 input_key_val: Some((key.clone(), val.map(StableRow))),
1429 }
1430 .as_plan(lir_id)
1431 }
1432 } else {
1433 let lir_id = self.allocate_lir_id();
1434 plan = LirRelationNode::Mfp {
1435 input: Box::new(plan),
1436 mfp: mfp_mir_to_lir_plan(mfp),
1437 input_key_val: input_key_val.map(|(key, val)| (key, val.map(StableRow))),
1438 }
1439 .as_plan(lir_id);
1440 keys = AvailableCollections::new_raw();
1441 }
1442 }
1443
1444 Ok(LoweredExpr {
1445 plan,
1446 keys,
1447 has_future_updates,
1448 })
1449 }
1450
1451 /// Lowers a `Reduce` with the given fields and an `mfp_on_top`, which is the MFP that is
1452 /// originally on top of the `Reduce`. This MFP, or a part of it, might be fused into the
1453 /// `Reduce`, in which case `mfp_on_top` is mutated to be the residual MFP, i.e., what was not
1454 /// fused.
1455 fn lower_reduce(
1456 &mut self,
1457 input: &MirRelationExpr,
1458 group_key: &Vec<MirScalarExpr>,
1459 aggregates: &Vec<AggregateExpr>,
1460 monotonic: &bool,
1461 expected_group_size: &Option<u64>,
1462 mfp_on_top: &mut MapFilterProject,
1463 fused_unnest_list: bool,
1464 ) -> Result<LoweredExpr, String> {
1465 let input_arity = input.arity();
1466 let LoweredExpr {
1467 plan: input,
1468 keys,
1469 has_future_updates: input_future,
1470 } = self.lower_mir_expr(input)?;
1471 let (input_key, permutation_and_new_arity) =
1472 if let Some((input_key, permutation, thinning)) = keys.arbitrary_arrangement() {
1473 (
1474 Some(input_key.clone()),
1475 Some((permutation.clone(), thinning.len() + input_key.len())),
1476 )
1477 } else {
1478 (None, None)
1479 };
1480 let key_val_plan = KeyValPlan::new(
1481 input_arity,
1482 group_key,
1483 aggregates,
1484 permutation_and_new_arity,
1485 );
1486 let mut reduce_plan = ReducePlan::create_from(
1487 aggregates.clone(),
1488 *monotonic,
1489 *expected_group_size,
1490 fused_unnest_list,
1491 );
1492
1493 // For single-time dataflows, upgrade a hierarchical reduce to its monotonic
1494 // variant with mandatory consolidation. `refine_single_time_consolidation`
1495 // later relaxes `must_consolidate` where the input is physically monotonic.
1496 // Selecting the variant before computing `keys` below keeps the advertised
1497 // `AvailableCollections` consistent with the final plan. `Reduce::keys()` is
1498 // the same for every hierarchical sub-variant, so the advertisement is in fact
1499 // identical either way.
1500 if self.single_time {
1501 if let ReducePlan::Hierarchical(hierarchical) = &mut reduce_plan {
1502 hierarchical.as_monotonic(true);
1503 }
1504 }
1505
1506 // Return the plan, and the keys it produces.
1507 let mfp_after;
1508 let output_arity;
1509 if self.enable_reduce_mfp_fusion {
1510 (mfp_after, *mfp_on_top, output_arity) =
1511 reduce_plan.extract_mfp_after(mfp_on_top.clone(), group_key.len());
1512 } else {
1513 (mfp_after, output_arity) = (
1514 MapFilterProject::new(mfp_on_top.input_arity),
1515 group_key.len() + aggregates.len(),
1516 );
1517 }
1518 soft_assert_eq_or_log!(
1519 mfp_on_top.input_arity,
1520 output_arity,
1521 "Output arity of reduce must match input arity for MFP on top of it"
1522 );
1523 let output_keys = reduce_plan.keys(group_key.len(), output_arity);
1524 let lir_id = self.allocate_lir_id();
1525 // `Reduce` builds its own input arrangement inside `render_reduce` (via `KeyValPlan`),
1526 // bypassing `ensure_collections`. So we can't piggy-back on an upstream `ArrangeBy`'s
1527 // strategy to request temporal bucketing on a temporal-MFP-fed input: there is no such
1528 // `ArrangeBy`. Instead we record the strategy directly on the `Reduce` node, and
1529 // `render_reduce` applies bucketing to the keyed `(key, val)` stream itself.
1530 let temporal_bucketing_strategy = strategy_from_future(input_future);
1531 // (This can't currently happen due to `extract_mfp_after` separating out any temporal part.)
1532 let has_future_updates = mfp_after.has_temporal_predicates();
1533 Ok(LoweredExpr {
1534 plan: LirRelationNode::Reduce {
1535 input_key,
1536 input: Box::new(input),
1537 key_val_plan,
1538 plan: reduce_plan,
1539 mfp_after: SafeMfpPlan::from_mfp(mfp_mir_to_lir(mfp_after)),
1540 temporal_bucketing_strategy,
1541 }
1542 .as_plan(lir_id),
1543 keys: output_keys,
1544 has_future_updates,
1545 })
1546 }
1547
1548 /// Replace the plan with another one
1549 /// that has the collection in some additional forms.
1550 pub fn arrange_by(
1551 &mut self,
1552 plan: LirRelationExpr,
1553 collections: AvailableCollections,
1554 old_collections: &AvailableCollections,
1555 arity: usize,
1556 has_future_updates: bool,
1557 ) -> LirRelationExpr {
1558 if let LirRelationExpr {
1559 node:
1560 LirRelationNode::ArrangeBy {
1561 input_key,
1562 input,
1563 input_mfp,
1564 mut forms,
1565 strategy,
1566 },
1567 lir_id,
1568 } = plan
1569 {
1570 forms.raw |= collections.raw;
1571 forms.arranged.extend(collections.arranged);
1572 forms.arranged.sort_by(|k1, k2| k1.0.cmp(&k2.0));
1573 forms.arranged.dedup_by(|k1, k2| k1.0 == k2.0);
1574 LirRelationNode::ArrangeBy {
1575 input_key,
1576 input,
1577 input_mfp,
1578 forms,
1579 strategy,
1580 }
1581 .as_plan(lir_id)
1582 } else {
1583 let (input_key, input_mfp) = if let Some((input_key, permutation, thinning)) =
1584 old_collections.arbitrary_arrangement()
1585 {
1586 let mut mfp = MapFilterProject::new(arity);
1587 mfp.permute_fn(|c| permutation[c], thinning.len() + input_key.len());
1588 (Some(input_key.clone()), mfp)
1589 } else {
1590 (None, MapFilterProject::new(arity))
1591 };
1592 let lir_id = self.allocate_lir_id();
1593
1594 LirRelationNode::ArrangeBy {
1595 input_key,
1596 input: Box::new(plan),
1597 input_mfp: mfp_mir_to_lir_plan(input_mfp),
1598 forms: collections,
1599 strategy: strategy_from_future(has_future_updates),
1600 }
1601 .as_plan(lir_id)
1602 }
1603 }
1604}
1605
1606/// Various bits of state to print along with error messages during LIR planning,
1607/// to aid debugging.
1608#[derive(Clone, Debug)]
1609pub struct LirDebugInfo {
1610 debug_name: String,
1611 id: GlobalId,
1612}
1613
1614impl std::fmt::Display for LirDebugInfo {
1615 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1616 write!(f, "Debug name: {}; id: {}", self.debug_name, self.id)
1617 }
1618}