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