Skip to main content

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}