Skip to main content

mz_compute_types/
plan.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//! An explicit representation of a rendering plan for provided dataflows.
11
12#![warn(missing_debug_implementations)]
13
14use std::collections::{BTreeMap, BTreeSet};
15
16use columnar::Columnar;
17use mz_expr::{
18    CollectionPlan, EvalError, Id, LetRecLimit, LocalId, MapFilterProject, MfpPlan, MirScalarExpr,
19    OptimizedMirRelationExpr, SafeMfpPlan, TableFunc,
20};
21use mz_ore::metric;
22use mz_ore::metrics::MetricsRegistry;
23use mz_ore::metrics::raw::IntCounterVec;
24use mz_ore::soft_assert_eq_no_log;
25use mz_ore::str::Indent;
26use mz_repr::explain::text::text_string_at;
27use mz_repr::explain::{DummyHumanizer, ExplainConfig, ExprHumanizer, PlanRenderingContext};
28use mz_repr::optimize::OptimizerFeatures;
29use mz_repr::{Diff, GlobalId, Row, Timestamp};
30use serde::{Deserialize, Serialize};
31
32use crate::dataflows::DataflowDescription;
33use crate::plan::join::JoinPlan;
34use crate::plan::reduce::{KeyValPlan, ReducePlan};
35use crate::plan::scalar::{LirScalarExpr, mfp_mir_to_lir_plan};
36use crate::plan::threshold::ThresholdPlan;
37use crate::plan::top_k::TopKPlan;
38use crate::plan::transform::{Transform, TransformConfig};
39
40mod lowering;
41
42pub mod interpret;
43pub mod join;
44pub mod reduce;
45pub mod render_plan;
46pub mod scalar;
47pub mod threshold;
48pub mod top_k;
49pub mod transform;
50
51/// Metrics collected during MIR to LIR lowering.
52#[derive(Debug, Clone)]
53pub struct LoweringMetrics {
54    /// Counts non-`None` results of `MapFilterProject::literal_constraints` during lowering,
55    /// labeled by the call site (`"get"` or `"mfp"`).
56    literal_constraints: IntCounterVec,
57}
58
59impl LoweringMetrics {
60    /// Registers the lowering metrics into `registry`.
61    pub fn register_into(registry: &MetricsRegistry) -> Self {
62        Self {
63            literal_constraints: registry.register(metric!(
64                name: "mz_optimizer_lowering_literal_constraints_total",
65                help: "How often the MFP-based literal-constraint detector succeeded, by call site.",
66                var_labels: ["case"],
67            )),
68        }
69    }
70
71    /// Records that a `literal_constraints` call at `case` produced a usable constraint.
72    pub fn inc_literal_constraints(&self, case: &str) {
73        self.literal_constraints.with_label_values(&[case]).inc();
74    }
75}
76
77/// The forms in which an operator's output is available.
78///
79/// These forms may include "raw", meaning as a streamed collection, but also any
80/// number of "arranged" representations.
81///
82/// Each arranged representation is described by a `KeyValRowMapping`, or rather
83/// at the moment by its three fields in a triple. These fields explain how to form
84/// a "key" by applying some expressions to each row, how to select "values" from
85/// columns not explicitly captured by the key, and how to return to the original
86/// row from the concatenation of key and value. Further explanation is available
87/// in the documentation for `KeyValRowMapping`.
88#[derive(
89    Clone,
90    Debug,
91    Default,
92    Deserialize,
93    Eq,
94    Ord,
95    PartialEq,
96    PartialOrd,
97    Serialize
98)]
99pub struct AvailableCollections {
100    /// Whether the collection exists in unarranged form.
101    pub raw: bool,
102    /// The list of available arrangements, presented as a `KeyValRowMapping`,
103    /// but here represented by a triple `(to_key, to_val, to_row)` instead.
104    /// The documentation for `KeyValRowMapping` explains these fields better.
105    pub arranged: Vec<(Vec<LirScalarExpr>, Vec<usize>, Vec<usize>)>,
106}
107
108impl AvailableCollections {
109    /// Represent a collection that has no arrangements.
110    pub fn new_raw() -> Self {
111        Self {
112            raw: true,
113            arranged: Vec::new(),
114        }
115    }
116
117    /// Represent a collection that is arranged in the specified ways.
118    pub fn new_arranged(arranged: Vec<(Vec<LirScalarExpr>, Vec<usize>, Vec<usize>)>) -> Self {
119        assert!(
120            !arranged.is_empty(),
121            "Invariant violated: at least one collection must exist"
122        );
123        Self {
124            raw: false,
125            arranged,
126        }
127    }
128
129    /// Get some arrangement, if one exists.
130    pub fn arbitrary_arrangement(&self) -> Option<&(Vec<LirScalarExpr>, Vec<usize>, Vec<usize>)> {
131        assert!(
132            self.raw || !self.arranged.is_empty(),
133            "Invariant violated: at least one collection must exist"
134        );
135        self.arranged.get(0)
136    }
137}
138
139/// How to render the arrangements requested by an `ArrangeBy`.
140///
141/// Decided during LIR lowering and consumed by the renderer. The variant says what the
142/// renderer will do, not what it knows about the input.
143#[derive(
144    Clone,
145    Copy,
146    Debug,
147    Deserialize,
148    Eq,
149    Ord,
150    PartialEq,
151    PartialOrd,
152    Serialize
153)]
154pub enum ArrangementStrategy {
155    /// Form arrangements directly from the input collection.
156    Direct,
157    /// Insert temporal bucketing in front of the arrangement, to delay future-stamped
158    /// updates (e.g., from `mz_now()` MFPs) until their bucket boundary releases them.
159    /// Honoured only when `ENABLE_COMPUTE_TEMPORAL_BUCKETING` is set; otherwise behaves like
160    /// `Direct`.
161    TemporalBucketing,
162}
163
164impl std::fmt::Display for ArrangementStrategy {
165    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
166        match self {
167            ArrangementStrategy::Direct => write!(f, "Direct"),
168            ArrangementStrategy::TemporalBucketing => write!(f, "TemporalBucketing"),
169        }
170    }
171}
172
173/// An identifier for an LIR node.
174#[derive(
175    Clone,
176    Copy,
177    Debug,
178    Deserialize,
179    Eq,
180    Ord,
181    PartialEq,
182    PartialOrd,
183    Serialize,
184    Columnar
185)]
186pub struct LirId(u64);
187
188impl LirId {
189    fn as_u64(&self) -> u64 {
190        self.0
191    }
192}
193
194impl From<LirId> for u64 {
195    fn from(value: LirId) -> Self {
196        value.as_u64()
197    }
198}
199
200impl std::fmt::Display for LirId {
201    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
202        write!(f, "{}", self.0)
203    }
204}
205
206/// A rendering plan with as much conditional logic as possible removed.
207#[derive(Clone, Debug, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
208pub struct LirRelationExpr {
209    /// A dataflow-local identifier.
210    pub lir_id: LirId,
211    /// The underlying operator.
212    pub node: LirRelationNode,
213}
214
215/// The actual AST node of the `LirRelationExpr`.
216#[derive(Clone, Debug, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
217pub enum LirRelationNode {
218    /// A collection containing a pre-determined collection.
219    Constant {
220        /// Explicit update triples for the collection.
221        rows: Result<Vec<(Row, Timestamp, Diff)>, EvalError>,
222    },
223    /// A reference to a bound collection.
224    ///
225    /// This is commonly either an external reference to an existing source or
226    /// maintained arrangement, or an internal reference to a `Let` identifier.
227    Get {
228        /// A global or local identifier naming the collection.
229        id: Id,
230        /// Arrangements that will be available.
231        ///
232        /// The collection will also be loaded if available, which it will
233        /// not be for imported data, but which it may be for locally defined
234        /// data.
235        // TODO: Be more explicit about whether a collection is available,
236        // although one can always produce it from an arrangement, and it
237        // seems generally advantageous to do that instead (to avoid cloning
238        // rows, by using `mfp` first on borrowed data).
239        keys: AvailableCollections,
240        /// The actions to take when introducing the collection.
241        plan: GetPlan,
242    },
243    /// Binds `value` to `id`, and then results in `body` with that binding.
244    ///
245    /// This stage has the effect of sharing `value` across multiple possible
246    /// uses in `body`, and is the only mechanism we have for sharing collection
247    /// information across parts of a dataflow.
248    ///
249    /// The binding is not available outside of `body`.
250    Let {
251        /// The local identifier to be used, available to `body` as `Id::Local(id)`.
252        id: LocalId,
253        /// The collection that should be bound to `id`.
254        value: Box<LirRelationExpr>,
255        /// The collection that results, which is allowed to contain `Get` stages
256        /// that reference `Id::Local(id)`.
257        body: Box<LirRelationExpr>,
258    },
259    /// Binds `values` to `ids`, evaluates them potentially recursively, and returns `body`.
260    ///
261    /// All bindings are available to all bindings, and to `body`.
262    /// The contents of each binding are initially empty, and then updated through a sequence
263    /// of iterations in which each binding is updated in sequence, from the most recent values
264    /// of all bindings.
265    LetRec {
266        /// The local identifiers to be used, available to `body` as `Id::Local(id)`.
267        ids: Vec<LocalId>,
268        /// The collection that should be bound to `id`.
269        values: Vec<LirRelationExpr>,
270        /// Maximum number of iterations. See further info on the MIR `LetRec`.
271        limits: Vec<Option<LetRecLimit>>,
272        /// The collection that results, which is allowed to contain `Get` stages
273        /// that reference `Id::Local(id)`.
274        body: Box<LirRelationExpr>,
275    },
276    /// Map, Filter, and Project operators.
277    ///
278    /// This stage contains work that we would ideally like to fuse to other plan
279    /// stages, but for practical reasons cannot. For example: threshold, topk,
280    /// and sometimes reduce stages are not able to absorb this operator.
281    Mfp {
282        /// The input collection.
283        input: Box<LirRelationExpr>,
284        /// Linear operator to apply to each record.
285        mfp: MfpPlan<LirScalarExpr>,
286        /// Whether the input is from an arrangement, and if so,
287        /// whether we can seek to a specific value therein
288        input_key_val: Option<(Vec<LirScalarExpr>, Option<Row>)>,
289    },
290    /// A variable number of output records for each input record.
291    ///
292    /// This stage is a bit of a catch-all for logic that does not easily fit in
293    /// map stages. This includes table valued functions, but also functions of
294    /// multiple arguments, and functions that modify the sign of updates.
295    ///
296    /// This stage allows a `MapFilterProject` operator to be fused to its output,
297    /// and this can be very important as otherwise the output of `func` is just
298    /// appended to the input record, for as many outputs as it has. This has the
299    /// unpleasant default behavior of repeating potentially large records that
300    /// are being unpacked, producing quadratic output in those cases. Instead,
301    /// in these cases use a `mfp` member that projects away these large fields.
302    FlatMap {
303        /// The particular arrangement of the input we expect to use,
304        /// if any
305        input_key: Option<Vec<LirScalarExpr>>,
306        /// The input collection.
307        input: Box<LirRelationExpr>,
308        /// Expressions that for each row prepare the arguments to `func`.
309        exprs: Vec<LirScalarExpr>,
310        /// The variable-record emitting function.
311        func: TableFunc,
312        /// Linear operator to apply to each record produced by `func`.
313        mfp_after: MfpPlan<LirScalarExpr>,
314    },
315    /// A multiway relational equijoin, with fused map, filter, and projection.
316    ///
317    /// This stage performs a multiway join among `inputs`, using the equality
318    /// constraints expressed in `plan`. The plan also describes the implementation
319    /// strategy we will use, and any pushed down per-record work.
320    Join {
321        /// An ordered list of inputs that will be joined.
322        inputs: Vec<LirRelationExpr>,
323        /// Detailed information about the implementation of the join.
324        ///
325        /// This includes information about the implementation strategy, but also
326        /// any map, filter, project work that we might follow the join with, but
327        /// potentially pushed down into the implementation of the join.
328        plan: JoinPlan,
329    },
330    /// Aggregation by key.
331    Reduce {
332        /// The particular arrangement of the input we expect to use,
333        /// if any
334        input_key: Option<Vec<LirScalarExpr>>,
335        /// The input collection.
336        input: Box<LirRelationExpr>,
337        /// A plan for changing input records into key, value pairs.
338        key_val_plan: KeyValPlan,
339        /// A plan for performing the reduce.
340        ///
341        /// The implementation of reduction has several different strategies based
342        /// on the properties of the reduction, and the input itself. Please check
343        /// out the documentation for this type for more detail.
344        plan: ReducePlan,
345        /// An MFP that must be applied to results. The projection part of this
346        /// MFP must preserve the key for the reduction; otherwise, the results
347        /// become undefined. Additionally, the MFP is guaranteed to be free from
348        /// temporal predicates so that it can be readily evaluated.
349        mfp_after: SafeMfpPlan<LirScalarExpr>,
350        /// Strategy for forming the internal input arrangement built by `Reduce`
351        /// (materialized via `key_val_plan`).
352        ///
353        /// Set by the lowering from the input's `has_future_updates` flag. The
354        /// renderer applies it to the keyed `(key, val)` stream feeding the
355        /// reduce. See `render_reduce` for the rationale on why this is
356        /// plumbed through `Reduce` rather than handled at the arrangement site.
357        ///
358        /// Note: unrelated to the hash buckets used by hierarchical reductions
359        /// (e.g. `ReducePlan::Hierarchical`'s `buckets`), which are an internal
360        /// sharding scheme for `min`/`max`-style aggregations. Here "bucketing"
361        /// refers exclusively to temporal (time-domain) bucketing of
362        /// future-stamped updates.
363        temporal_bucketing_strategy: ArrangementStrategy,
364    },
365    /// Key-based "Top K" operator, retaining the first K records in each group.
366    TopK {
367        /// The input collection.
368        input: Box<LirRelationExpr>,
369        /// A plan for performing the Top-K.
370        ///
371        /// The implementation of reduction has several different strategies based
372        /// on the properties of the reduction, and the input itself. Please check
373        /// out the documentation for this type for more detail.
374        top_k_plan: TopKPlan,
375        /// Strategy for bucketing the input collection ahead of the Top-K operator.
376        ///
377        /// Set by the lowering from the input's `has_future_updates` flag. The
378        /// renderer applies it to the per-row input stream at the top of
379        /// `render_topk`, covering all three `TopKPlan` arms uniformly. See
380        /// `LirRelationNode::Reduce::temporal_bucketing_strategy` for the underlying
381        /// convention.
382        temporal_bucketing_strategy: ArrangementStrategy,
383    },
384    /// Inverts the sign of each update.
385    Negate {
386        /// The input collection.
387        input: Box<LirRelationExpr>,
388    },
389    /// Filters records that accumulate negatively.
390    ///
391    /// Although the operator suppresses updates, it is a stateful operator taking
392    /// resources proportional to the number of records with non-zero accumulation.
393    Threshold {
394        /// The input collection.
395        input: Box<LirRelationExpr>,
396        /// A plan for performing the threshold.
397        ///
398        /// The implementation of reduction has several different strategies based
399        /// on the properties of the reduction, and the input itself. Please check
400        /// out the documentation for this type for more detail.
401        threshold_plan: ThresholdPlan,
402    },
403    /// Adds the contents of the input collections.
404    ///
405    /// Importantly, this is *multiset* union, so the multiplicities of records will
406    /// add. This is in contrast to *set* union, where the multiplicities would be
407    /// capped at one. A set union can be formed with `Union` followed by `Reduce`
408    /// implementing the "distinct" operator.
409    Union {
410        /// The input collections
411        inputs: Vec<LirRelationExpr>,
412        /// Whether to consolidate the output, e.g., cancel negated records.
413        consolidate_output: bool,
414        /// Per-input bucketing strategies. Lockstep with `inputs`: index `i` is the
415        /// strategy applied to `inputs[i]` before concatenation.
416        ///
417        /// Set by the lowering from each input's `has_future_updates` flag. Only
418        /// consolidating Unions (`consolidate_output: true`) carry non-`Direct`
419        /// entries, because bucketing only pays off ahead of a consolidating
420        /// downstream operator. See `LirRelationNode::Reduce::temporal_bucketing_strategy`
421        /// for the underlying convention.
422        temporal_bucketing_strategies: Vec<ArrangementStrategy>,
423    },
424    /// The `input` plan, but with additional arrangements.
425    ///
426    /// This operator does not change the logical contents of `input`, but ensures
427    /// that certain arrangements are available in the results. This operator can
428    /// be important for e.g. the `Join` stage which benefits from multiple arrangements
429    /// or to cap a `LirRelationExpr` so that indexes can be exported.
430    ArrangeBy {
431        /// The key that must be used to access the input.
432        input_key: Option<Vec<LirScalarExpr>>,
433        /// The input collection.
434        input: Box<LirRelationExpr>,
435        /// The MFP that must be applied to the input.
436        input_mfp: MfpPlan<LirScalarExpr>,
437        /// A list of arrangement keys, and possibly a raw collection,
438        /// that will be added to those of the input. Does not include
439        /// any other existing arrangements.
440        forms: AvailableCollections,
441        /// How the renderer should form the arrangements requested by `forms`.
442        strategy: ArrangementStrategy,
443    },
444}
445
446impl LirRelationNode {
447    /// Iterates through references to child expressions.
448    pub fn children(&self) -> impl Iterator<Item = &LirRelationExpr> {
449        let mut first = None;
450        let mut second = None;
451        let mut rest = None;
452        let mut last = None;
453
454        use LirRelationNode::*;
455        match self {
456            Constant { .. } | Get { .. } => (),
457            Let { value, body, .. } => {
458                first = Some(&**value);
459                second = Some(&**body);
460            }
461            LetRec { values, body, .. } => {
462                rest = Some(values);
463                last = Some(&**body);
464            }
465            Mfp { input, .. }
466            | FlatMap { input, .. }
467            | Reduce { input, .. }
468            | TopK { input, .. }
469            | Negate { input, .. }
470            | Threshold { input, .. }
471            | ArrangeBy { input, .. } => {
472                first = Some(&**input);
473            }
474            Join { inputs, .. } | Union { inputs, .. } => {
475                rest = Some(inputs);
476            }
477        }
478
479        first
480            .into_iter()
481            .chain(second)
482            .chain(rest.into_iter().flatten())
483            .chain(last)
484    }
485
486    /// Iterates through mutable references to child expressions.
487    pub fn children_mut(&mut self) -> impl Iterator<Item = &mut LirRelationExpr> {
488        let mut first = None;
489        let mut second = None;
490        let mut rest = None;
491        let mut last = None;
492
493        use LirRelationNode::*;
494        match self {
495            Constant { .. } | Get { .. } => (),
496            Let { value, body, .. } => {
497                first = Some(&mut **value);
498                second = Some(&mut **body);
499            }
500            LetRec { values, body, .. } => {
501                rest = Some(values);
502                last = Some(&mut **body);
503            }
504            Mfp { input, .. }
505            | FlatMap { input, .. }
506            | Reduce { input, .. }
507            | TopK { input, .. }
508            | Negate { input, .. }
509            | Threshold { input, .. }
510            | ArrangeBy { input, .. } => {
511                first = Some(&mut **input);
512            }
513            Join { inputs, .. } | Union { inputs, .. } => {
514                rest = Some(inputs);
515            }
516        }
517
518        first
519            .into_iter()
520            .chain(second)
521            .chain(rest.into_iter().flatten())
522            .chain(last)
523    }
524}
525
526impl LirRelationNode {
527    /// Attach an `lir_id` to a `LirRelationNode` to make a complete `LirRelationExpr`.
528    pub fn as_plan(self, lir_id: LirId) -> LirRelationExpr {
529        LirRelationExpr { lir_id, node: self }
530    }
531}
532
533impl LirRelationExpr {
534    /// Pretty-print this [LirRelationExpr] to a string.
535    pub fn pretty(&self) -> String {
536        let config = ExplainConfig::default();
537        self.debug_explain(&config, None)
538    }
539
540    /// Pretty-print this [LirRelationExpr] to a string using a custom
541    /// [ExplainConfig] and an optionally provided [ExprHumanizer].
542    /// This is intended for debugging and tests, not users.
543    pub fn debug_explain(
544        &self,
545        config: &ExplainConfig,
546        humanizer: Option<&dyn ExprHumanizer>,
547    ) -> String {
548        text_string_at(self, || PlanRenderingContext {
549            indent: Indent::default(),
550            humanizer: humanizer.unwrap_or(&DummyHumanizer),
551            annotations: BTreeMap::default(),
552            config,
553            ambiguous_ids: BTreeSet::default(),
554        })
555    }
556}
557
558/// How a `Get` stage will be rendered.
559#[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq, Ord, PartialOrd)]
560pub enum GetPlan {
561    /// Simply pass input arrangements on to the next stage.
562    PassArrangements,
563    /// Using the supplied key, optionally seek the row, and apply the MFP.
564    Arrangement(Vec<LirScalarExpr>, Option<Row>, MfpPlan<LirScalarExpr>),
565    /// Scan the input collection (unarranged) and apply the MFP.
566    Collection(MfpPlan<LirScalarExpr>),
567}
568
569impl LirRelationExpr {
570    /// Convert the dataflow description into one that uses render plans.
571    #[mz_ore::instrument(
572        target = "optimizer",
573        level = "debug",
574        fields(path.segment = "finalize_dataflow")
575    )]
576    pub fn finalize_dataflow(
577        desc: DataflowDescription<OptimizedMirRelationExpr>,
578        features: &OptimizerFeatures,
579        metrics: Option<&LoweringMetrics>,
580    ) -> Result<DataflowDescription<Self>, String> {
581        // First, we lower the dataflow description from MIR to LIR.
582        let mut dataflow = Self::lower_dataflow(desc, features, metrics)?;
583
584        // Subsequently, we perform plan refinements for the dataflow.
585        Self::refine_source_mfps(&mut dataflow);
586
587        // Note: `consolidate_output` for `Union` and per-input
588        // `temporal_bucketing_strategies` are decided at lowering time (see the
589        // `Union` arm of `lower_mir_expr_stack_safe`). The pre-existing
590        // `refine_union_negate_consolidation` pass — which used to flip
591        // `consolidate_output` to `true` for Unions with a `Negate` child — has
592        // been folded into the lowering, since lowering is the only point where
593        // the bucketing decision (which depends on `has_future_updates`) is
594        // available.
595
596        if dataflow.is_single_time() {
597            // The relaxation of the `must_consolidate` flag performs an LIR-based
598            // analysis and transform under checked recursion. By a similar argument
599            // made in `from_mir`, we do not expect the recursion limit to be hit.
600            // However, if that happens, we propagate an error to the caller.
601            // To apply the transform, we first obtain monotonic source and index
602            // global IDs and add them to a `TransformConfig` instance.
603            let monotonic_ids = dataflow
604                .source_imports
605                .iter()
606                .filter_map(|(id, source_import)| source_import.monotonic.then_some(*id))
607                .chain(
608                    dataflow
609                        .index_imports
610                        .iter()
611                        .filter_map(|(_id, index_import)| {
612                            if index_import.monotonic {
613                                Some(index_import.desc.on_id)
614                            } else {
615                                None
616                            }
617                        }),
618                )
619                .collect::<BTreeSet<_>>();
620
621            let config = TransformConfig { monotonic_ids };
622            Self::refine_single_time_consolidation(&mut dataflow, &config)?;
623        }
624
625        soft_assert_eq_no_log!(dataflow.check_invariants(), Ok(()));
626
627        mz_repr::explain::trace_plan(&dataflow);
628
629        Ok(dataflow)
630    }
631
632    /// Lowers the dataflow description from MIR to LIR. To this end, the
633    /// method collects all available arrangements and based on this information
634    /// creates plans for every object to be built for the dataflow.
635    #[mz_ore::instrument(
636        target = "optimizer",
637        level = "debug",
638        fields(path.segment ="mir_to_lir")
639    )]
640    fn lower_dataflow(
641        desc: DataflowDescription<OptimizedMirRelationExpr>,
642        features: &OptimizerFeatures,
643        metrics: Option<&LoweringMetrics>,
644    ) -> Result<DataflowDescription<Self>, String> {
645        let context = lowering::Context::new(desc.debug_name.clone(), features, metrics);
646        let dataflow = context.lower(desc)?;
647
648        mz_repr::explain::trace_plan(&dataflow);
649
650        Ok(dataflow)
651    }
652
653    /// Refines the source instance descriptions for sources imported by `dataflow` to
654    /// push down common MFP expressions.
655    #[mz_ore::instrument(
656        target = "optimizer",
657        level = "debug",
658        fields(path.segment = "refine_source_mfps")
659    )]
660    fn refine_source_mfps(dataflow: &mut DataflowDescription<Self>) {
661        use crate::plan::scalar::mfp_plan_lir_to_mir;
662
663        for (source_id, source_import) in dataflow.source_imports.iter_mut() {
664            let source = &mut source_import.desc;
665            let source_id = *source_id;
666            let mut identity_present = false;
667
668            // First pass: swap MfpPlans out of GetPlan::Collection nodes,
669            // recording their LirId so we can put them back.
670            let mut taken: Vec<(LirId, MfpPlan<LirScalarExpr>)> = Vec::new();
671            for build_desc in dataflow.objects_to_build.iter_mut() {
672                let mut todo = vec![&mut build_desc.plan];
673                while let Some(expression) = todo.pop() {
674                    let lir_id = expression.lir_id;
675                    let node = &mut expression.node;
676                    if let LirRelationNode::Get { id, plan, .. } = node {
677                        if *id == mz_expr::Id::Global(source_id) {
678                            match plan {
679                                GetPlan::Collection(mfp_plan) => {
680                                    let arity = mfp_plan.safe_mfp().projection.len();
681                                    let placeholder = MfpPlan::from_parts(
682                                        mz_expr::SafeMfpPlan::from_mfp(MapFilterProject::new(
683                                            arity,
684                                        )),
685                                        Vec::new(),
686                                        Vec::new(),
687                                    );
688                                    taken.push((lir_id, std::mem::replace(mfp_plan, placeholder)));
689                                }
690                                GetPlan::PassArrangements => {
691                                    identity_present = true;
692                                }
693                                GetPlan::Arrangement(..) => {
694                                    panic!("Surprising `GetPlan` for imported source: {:?}", plan);
695                                }
696                            }
697                        }
698                    } else {
699                        todo.extend(node.children_mut());
700                    }
701                }
702            }
703
704            // Direct exports of sources are possible, and prevent pushdown.
705            identity_present |= dataflow
706                .index_exports
707                .values()
708                .any(|(x, _)| x.on_id == source_id);
709            identity_present |= dataflow.sink_exports.values().any(|x| x.from == source_id);
710
711            // Build a map from LirId → new MfpPlan to put back.
712            let replacements: BTreeMap<LirId, MfpPlan<LirScalarExpr>> =
713                if !identity_present && !taken.is_empty() {
714                    // Convert LIR MfpPlans → MIR MapFilterProjects by folding
715                    // temporal bounds back as mz_now() predicates, so that
716                    // extract_common's column remapping applies uniformly.
717                    let mut mir_mfps: Vec<(LirId, MapFilterProject<MirScalarExpr>)> = taken
718                        .into_iter()
719                        .map(|(lir_id, lir_plan)| {
720                            let mir_mfp = mfp_plan_lir_to_mir(lir_plan).into_map_filter_project();
721                            (lir_id, mir_mfp)
722                        })
723                        .collect();
724                    let mut mfp_refs: Vec<&mut MapFilterProject<MirScalarExpr>> =
725                        mir_mfps.iter_mut().map(|(_, mfp)| mfp).collect();
726
727                    let common = MapFilterProject::extract_common(&mut mfp_refs[..]);
728                    let mut source_mfp = if let Some(mfp) = source.arguments.operators.take() {
729                        MapFilterProject::compose(mfp, common)
730                    } else {
731                        common
732                    };
733                    source_mfp.optimize();
734                    source.arguments.operators = Some(source_mfp);
735
736                    // Convert mutated MIR MFPs back to LIR MfpPlans.
737                    mir_mfps
738                        .into_iter()
739                        .map(|(lir_id, mir_mfp)| (lir_id, mfp_mir_to_lir_plan(mir_mfp)))
740                        .collect()
741                } else {
742                    taken.into_iter().collect()
743                };
744
745            // Second pass: put the MfpPlans back by LirId.
746            for build_desc in dataflow.objects_to_build.iter_mut() {
747                let mut todo = vec![&mut build_desc.plan];
748                while let Some(expression) = todo.pop() {
749                    if let Some(replacement) = replacements.get(&expression.lir_id) {
750                        if let LirRelationNode::Get {
751                            plan: GetPlan::Collection(mfp_plan),
752                            ..
753                        } = &mut expression.node
754                        {
755                            *mfp_plan = replacement.clone();
756                        } else {
757                            panic!(
758                                "LirId {:?} was a GetPlan::Collection but is now {:?}",
759                                expression.lir_id, expression.node
760                            );
761                        }
762                    }
763                    todo.extend(expression.node.children_mut());
764                }
765            }
766        }
767        mz_repr::explain::trace_plan(dataflow);
768    }
769
770    /// Refines the plans of objects to be built as part of a single-time `dataflow` to relax
771    /// the setting of the `must_consolidate` attribute of monotonic operators, if necessary,
772    /// whenever the input is deemed to be physically monotonic.
773    #[mz_ore::instrument(
774        target = "optimizer",
775        level = "debug",
776        fields(path.segment = "refine_single_time_consolidation")
777    )]
778    fn refine_single_time_consolidation(
779        dataflow: &mut DataflowDescription<Self>,
780        config: &TransformConfig,
781    ) -> Result<(), String> {
782        // We should only reach here if we have a one-shot SELECT query, i.e.,
783        // a single-time dataflow.
784        assert!(dataflow.is_single_time());
785
786        let transform = transform::RelaxMustConsolidate;
787        for build_desc in dataflow.objects_to_build.iter_mut() {
788            transform
789                .transform(config, &mut build_desc.plan)
790                .map_err(|_| "Maximum recursion limit error in consolidation relaxation.")?;
791        }
792        mz_repr::explain::trace_plan(dataflow);
793        Ok(())
794    }
795}
796
797impl CollectionPlan for LirRelationNode {
798    fn depends_on_into(&self, out: &mut BTreeSet<GlobalId>) {
799        match self {
800            LirRelationNode::Constant { rows: _ } => (),
801            LirRelationNode::Get {
802                id,
803                keys: _,
804                plan: _,
805            } => match id {
806                Id::Global(id) => {
807                    out.insert(*id);
808                }
809                Id::Local(_) => (),
810            },
811            LirRelationNode::Let { id: _, value, body } => {
812                value.depends_on_into(out);
813                body.depends_on_into(out);
814            }
815            LirRelationNode::LetRec {
816                ids: _,
817                values,
818                limits: _,
819                body,
820            } => {
821                for value in values.iter() {
822                    value.depends_on_into(out);
823                }
824                body.depends_on_into(out);
825            }
826            LirRelationNode::Join { inputs, plan: _ }
827            | LirRelationNode::Union {
828                inputs,
829                consolidate_output: _,
830                temporal_bucketing_strategies: _,
831            } => {
832                for input in inputs {
833                    input.depends_on_into(out);
834                }
835            }
836            LirRelationNode::Mfp {
837                input,
838                mfp: _,
839                input_key_val: _,
840            }
841            | LirRelationNode::FlatMap {
842                input_key: _,
843                input,
844                exprs: _,
845                func: _,
846                mfp_after: _,
847            }
848            | LirRelationNode::ArrangeBy {
849                input_key: _,
850                input,
851                input_mfp: _,
852                forms: _,
853                strategy: _,
854            }
855            | LirRelationNode::Reduce {
856                input_key: _,
857                input,
858                key_val_plan: _,
859                plan: _,
860                mfp_after: _,
861                temporal_bucketing_strategy: _,
862            }
863            | LirRelationNode::TopK {
864                input,
865                top_k_plan: _,
866                temporal_bucketing_strategy: _,
867            }
868            | LirRelationNode::Negate { input }
869            | LirRelationNode::Threshold {
870                input,
871                threshold_plan: _,
872            } => {
873                input.depends_on_into(out);
874            }
875        }
876    }
877}
878
879impl CollectionPlan for LirRelationExpr {
880    fn depends_on_into(&self, out: &mut BTreeSet<GlobalId>) {
881        self.node.depends_on_into(out);
882    }
883}
884
885/// Returns bucket sizes, descending, suitable for hierarchical decomposition of an operator, based
886/// on the expected number of rows that will have the same group key.
887fn bucketing_of_expected_group_size(expected_group_size: Option<u64>) -> Vec<u64> {
888    // NOTE(vmarcos): The fan-in of 16 defined below is used in the tuning advice built-in view
889    // mz_introspection.mz_expected_group_size_advice.
890    let mut buckets = vec![];
891    let mut current = 16;
892
893    // Plan for 4B records in the expected case if the user didn't specify a group size.
894    let limit = expected_group_size.unwrap_or(4_000_000_000);
895
896    // Distribute buckets in powers of 16, so that we can strike a balance between how many inputs
897    // each layer gets from the preceding layer, while also limiting the number of layers.
898    while current < limit {
899        buckets.push(current);
900        current = current.saturating_mul(16);
901    }
902
903    buckets.reverse();
904    buckets
905}