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